add test for agent CLI

coder · sreya · May 18, 2023 · May 17, 2023 · May 17, 2023 · May 17, 2023
commit 2b4140f39cac809276867ccecfe8e43e67a473ff
diff --git a/agent/agent.go b/agent/agent.go
@@ -62,7 +62,7 @@ type Options struct {
 	IgnorePorts            map[int]string
 	SSHMaxTimeout          time.Duration
 	TailnetListenPort      uint16
-	Subsystem              agentsdk.AgentSubsystem
+	Subsystem              codersdk.AgentSubsystem
 }
 
 type Client interface {
@@ -138,7 +138,7 @@ type agent struct {
 	// listing all listening ports. This is helpful to hide ports that
 	// are used by the agent, that the user does not care about.
 	ignorePorts map[int]string
-	subsystem   agentsdk.AgentSubsystem
+	subsystem   codersdk.AgentSubsystem
 
 	reconnectingPTYs       sync.Map
 	reconnectingPTYTimeout time.Duration
@@ -1467,11 +1467,11 @@ const EnvAgentSubsystem = "CODER_AGENT_SUBSYSTEM"
 
 // SubsystemFromEnv returns the subsystem (if any) the agent
 // is running inside of.
-func SubsystemFromEnv() agentsdk.AgentSubsystem {
+func SubsystemFromEnv() codersdk.AgentSubsystem {
 	ss := os.Getenv(EnvAgentSubsystem)
-	switch agentsdk.AgentSubsystem(ss) {
-	case agentsdk.AgentSubsystemEnvbox:
-		return agentsdk.AgentSubsystemEnvbox
+	switch codersdk.AgentSubsystem(ss) {
+	case codersdk.AgentSubsystemEnvbox:
+		return codersdk.AgentSubsystemEnvbox
 	default:
 		return ""
 	}

diff --git a/cli/agent.go b/cli/agent.go
@@ -26,6 +26,7 @@ import (
 	"github.com/coder/coder/agent/reaper"
 	"github.com/coder/coder/buildinfo"
 	"github.com/coder/coder/cli/clibase"
+	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
 )
 
@@ -197,6 +198,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				return xerrors.Errorf("add executable to $PATH: %w", err)
 			}
 
+			subsystem := inv.Environ.Get(agent.EnvAgentSubsystem)
 			agnt := agent.New(agent.Options{
 				Client:            client,
 				Logger:            logger,
@@ -218,7 +220,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				},
 				IgnorePorts:   ignorePorts,
 				SSHMaxTimeout: sshMaxTimeout,
-				Subsystem:     agent.SubsystemFromEnv(),
+				Subsystem:     codersdk.AgentSubsystem(subsystem),
 			})
 
 			debugSrvClose := ServeHandler(ctx, logger, agnt.HTTPDebug(), debugAddress, "debug")

diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -12,8 +12,10 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/coder/coder/agent"
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
 	"github.com/coder/coder/pty/ptytest"
@@ -235,4 +237,43 @@ func TestWorkspaceAgent(t *testing.T) {
 		_, err = uuid.Parse(strings.TrimSpace(string(token)))
 		require.NoError(t, err)
 	})
+
+	t.Run("PostStartup", func(t *testing.T) {
+		t.Parallel()
+
+		authToken := uuid.NewString()
+		client := coderdtest.New(t, &coderdtest.Options{
+			IncludeProvisionerDaemon: true,
+		})
+		user := coderdtest.CreateFirstUser(t, client)
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse:          echo.ParseComplete,
+			ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
+		})
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+
+		logDir := t.TempDir()
+		inv, _ := clitest.New(t,
+			"agent",
+			"--auth", "token",
+			"--agent-token", authToken,
+			"--agent-url", client.URL.String(),
+			"--log-dir", logDir,
+		)
+		// Set the subsystem for the agent.
+		inv.Environ.Set(agent.EnvAgentSubsystem, string(codersdk.AgentSubsystemEnvbox))
+
+		pty := ptytest.New(t).Attach(inv)
+
+		clitest.Start(t, inv)
+		pty.ExpectMatch("starting agent")
+
+		resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+		require.Len(t, resources, 1)
+		require.Len(t, resources[0].Agents, 1)
+		require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystem)
+	})
 }
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -1130,6 +1130,7 @@ func convertWorkspaceAgent(derpMap *tailcfg.DERPMap, coordinator tailnet.Coordin
 		StartupScriptTimeoutSeconds:  dbAgent.StartupScriptTimeoutSeconds,
 		ShutdownScript:               dbAgent.ShutdownScript.String,
 		ShutdownScriptTimeoutSeconds: dbAgent.ShutdownScriptTimeoutSeconds,
+		Subsystem:                    codersdk.AgentSubsystem(dbAgent.Subsystem),
 	}
 	node := coordinator.Node(dbAgent.ID)
 	if node != nil {
@@ -1985,9 +1986,9 @@ func convertWorkspaceAgentStartupLog(logEntry database.WorkspaceAgentStartupLog)
 	}
 }
 
-func convertWorkspaceAgentSubsystem(ss agentsdk.AgentSubsystem) database.WorkspaceAgentSubsystem {
+func convertWorkspaceAgentSubsystem(ss codersdk.AgentSubsystem) database.WorkspaceAgentSubsystem {
 	switch ss {
-	case agentsdk.AgentSubsystemEnvbox:
+	case codersdk.AgentSubsystemEnvbox:
 		return database.WorkspaceAgentSubsystemEnvbox
 	default:
 		return database.WorkspaceAgentSubsystemNone

diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -453,12 +453,6 @@ func (c *Client) ReportStats(ctx context.Context, log slog.Logger, statsChan <-c
 	}), nil
 }
 
-type AgentSubsystem string
-
-const (
-	AgentSubsystemEnvbox AgentSubsystem = "envbox"
-)
-
 // Stats records the Agent's network connection statistics for use in
 // user-facing metrics and debugging.
 type Stats struct {
@@ -550,9 +544,9 @@ func (c *Client) PostLifecycle(ctx context.Context, req PostLifecycleRequest) er
 }
 
 type PostStartupRequest struct {
-	Version           string         `json:"version"`
-	ExpandedDirectory string         `json:"expanded_directory"`
-	Subsystem         AgentSubsystem `json:"subsystem"`
+	Version           string                  `json:"version"`
+	ExpandedDirectory string                  `json:"expanded_directory"`
+	Subsystem         codersdk.AgentSubsystem `json:"subsystem"`
 }
 
 func (c *Client) PostStartup(ctx context.Context, req PostStartupRequest) error {

diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -130,9 +130,10 @@ type WorkspaceAgent struct {
 	// LoginBeforeReady if true, the agent will delay logins until it is ready (e.g. executing startup script has ended).
 	LoginBeforeReady bool `json:"login_before_ready"`
 	// StartupScriptTimeoutSeconds is the number of seconds to wait for the startup script to complete. If the script does not complete within this time, the agent lifecycle will be marked as start_timeout.
-	StartupScriptTimeoutSeconds  int32  `json:"startup_script_timeout_seconds"`
-	ShutdownScript               string `json:"shutdown_script,omitempty"`
-	ShutdownScriptTimeoutSeconds int32  `json:"shutdown_script_timeout_seconds"`
+	StartupScriptTimeoutSeconds  int32          `json:"startup_script_timeout_seconds"`
+	ShutdownScript               string         `json:"shutdown_script,omitempty"`
+	ShutdownScriptTimeoutSeconds int32          `json:"shutdown_script_timeout_seconds"`
+	Subsystem                    AgentSubsystem `json:"subsytem"`
 }
 
 type DERPRegion struct {
@@ -553,3 +554,9 @@ type WorkspaceAgentStartupLog struct {
 	Output    string    `json:"output"`
 	Level     LogLevel  `json:"level"`
 }
+
+type AgentSubsystem string
+
+const (
+	AgentSubsystemEnvbox AgentSubsystem = "envbox"
+)