Skip to content

Commit 2b4140f

Browse files
sreyasreya
committed
add test for agent CLI
1 parent 7ef6911 commit 2b4140f

File tree

6 files changed

+66
-21
lines changed

6 files changed

+66
-21
lines changed

‎agent/agent.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ type Options struct {
6262
IgnorePorts map[int]string
6363
SSHMaxTimeout time.Duration
6464
TailnetListenPort uint16
65-
Subsystem agentsdk.AgentSubsystem
65+
Subsystem codersdk.AgentSubsystem
6666
}
6767

6868
type Client interface {
@@ -138,7 +138,7 @@ type agent struct {
138138
// listing all listening ports. This is helpful to hide ports that
139139
// are used by the agent, that the user does not care about.
140140
ignorePorts map[int]string
141-
subsystem agentsdk.AgentSubsystem
141+
subsystem codersdk.AgentSubsystem
142142

143143
reconnectingPTYs sync.Map
144144
reconnectingPTYTimeout time.Duration
@@ -1467,11 +1467,11 @@ const EnvAgentSubsystem = "CODER_AGENT_SUBSYSTEM"
14671467

14681468
// SubsystemFromEnv returns the subsystem (if any) the agent
14691469
// is running inside of.
1470-
func SubsystemFromEnv() agentsdk.AgentSubsystem {
1470+
func SubsystemFromEnv() codersdk.AgentSubsystem {
14711471
ss := os.Getenv(EnvAgentSubsystem)
1472-
switch agentsdk.AgentSubsystem(ss) {
1473-
case agentsdk.AgentSubsystemEnvbox:
1474-
return agentsdk.AgentSubsystemEnvbox
1472+
switch codersdk.AgentSubsystem(ss) {
1473+
case codersdk.AgentSubsystemEnvbox:
1474+
return codersdk.AgentSubsystemEnvbox
14751475
default:
14761476
return ""
14771477
}

‎cli/agent.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
"github.com/coder/coder/agent/reaper"
2727
"github.com/coder/coder/buildinfo"
2828
"github.com/coder/coder/cli/clibase"
29+
"github.com/coder/coder/codersdk"
2930
"github.com/coder/coder/codersdk/agentsdk"
3031
)
3132

@@ -197,6 +198,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
197198
return xerrors.Errorf("add executable to $PATH: %w", err)
198199
}
199200

201+
subsystem := inv.Environ.Get(agent.EnvAgentSubsystem)
200202
agnt := agent.New(agent.Options{
201203
Client: client,
202204
Logger: logger,
@@ -218,7 +220,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
218220
},
219221
IgnorePorts: ignorePorts,
220222
SSHMaxTimeout: sshMaxTimeout,
221-
Subsystem: agent.SubsystemFromEnv(),
223+
Subsystem: codersdk.AgentSubsystem(subsystem),
222224
})
223225

224226
debugSrvClose := ServeHandler(ctx, logger, agnt.HTTPDebug(), debugAddress, "debug")

‎cli/agent_test.go

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,10 @@ import (
1212
"github.com/stretchr/testify/assert"
1313
"github.com/stretchr/testify/require"
1414

15+
"github.com/coder/coder/agent"
1516
"github.com/coder/coder/cli/clitest"
1617
"github.com/coder/coder/coderd/coderdtest"
18+
"github.com/coder/coder/codersdk"
1719
"github.com/coder/coder/provisioner/echo"
1820
"github.com/coder/coder/provisionersdk/proto"
1921
"github.com/coder/coder/pty/ptytest"
@@ -235,4 +237,43 @@ func TestWorkspaceAgent(t *testing.T) {
235237
_, err = uuid.Parse(strings.TrimSpace(string(token)))
236238
require.NoError(t, err)
237239
})
240+
241+
t.Run("PostStartup", func(t *testing.T) {
242+
t.Parallel()
243+
244+
authToken := uuid.NewString()
245+
client := coderdtest.New(t, &coderdtest.Options{
246+
IncludeProvisionerDaemon: true,
247+
})
248+
user := coderdtest.CreateFirstUser(t, client)
249+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
250+
Parse: echo.ParseComplete,
251+
ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
252+
})
253+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
254+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
255+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
256+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
257+
258+
logDir := t.TempDir()
259+
inv, _ := clitest.New(t,
260+
"agent",
261+
"--auth", "token",
262+
"--agent-token", authToken,
263+
"--agent-url", client.URL.String(),
264+
"--log-dir", logDir,
265+
)
266+
// Set the subsystem for the agent.
267+
inv.Environ.Set(agent.EnvAgentSubsystem, string(codersdk.AgentSubsystemEnvbox))
268+
269+
pty := ptytest.New(t).Attach(inv)
270+
271+
clitest.Start(t, inv)
272+
pty.ExpectMatch("starting agent")
273+
274+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
275+
require.Len(t, resources, 1)
276+
require.Len(t, resources[0].Agents, 1)
277+
require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystem)
278+
})
238279
}

‎coderd/workspaceagents.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1130,6 +1130,7 @@ func convertWorkspaceAgent(derpMap *tailcfg.DERPMap, coordinator tailnet.Coordin
11301130
StartupScriptTimeoutSeconds: dbAgent.StartupScriptTimeoutSeconds,
11311131
ShutdownScript: dbAgent.ShutdownScript.String,
11321132
ShutdownScriptTimeoutSeconds: dbAgent.ShutdownScriptTimeoutSeconds,
1133+
Subsystem: codersdk.AgentSubsystem(dbAgent.Subsystem),
11331134
}
11341135
node := coordinator.Node(dbAgent.ID)
11351136
if node != nil {
@@ -1985,9 +1986,9 @@ func convertWorkspaceAgentStartupLog(logEntry database.WorkspaceAgentStartupLog)
19851986
}
19861987
}
19871988

1988-
func convertWorkspaceAgentSubsystem(ss agentsdk.AgentSubsystem) database.WorkspaceAgentSubsystem {
1989+
func convertWorkspaceAgentSubsystem(ss codersdk.AgentSubsystem) database.WorkspaceAgentSubsystem {
19891990
switch ss {
1990-
case agentsdk.AgentSubsystemEnvbox:
1991+
case codersdk.AgentSubsystemEnvbox:
19911992
return database.WorkspaceAgentSubsystemEnvbox
19921993
default:
19931994
return database.WorkspaceAgentSubsystemNone

‎codersdk/agentsdk/agentsdk.go

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -453,12 +453,6 @@ func (c *Client) ReportStats(ctx context.Context, log slog.Logger, statsChan <-c
453453
}), nil
454454
}
455455

456-
type AgentSubsystem string
457-
458-
const (
459-
AgentSubsystemEnvbox AgentSubsystem = "envbox"
460-
)
461-
462456
// Stats records the Agent's network connection statistics for use in
463457
// user-facing metrics and debugging.
464458
type Stats struct {
@@ -550,9 +544,9 @@ func (c *Client) PostLifecycle(ctx context.Context, req PostLifecycleRequest) er
550544
}
551545

552546
type PostStartupRequest struct {
553-
Version string `json:"version"`
554-
ExpandedDirectory string `json:"expanded_directory"`
555-
Subsystem AgentSubsystem `json:"subsystem"`
547+
Version string `json:"version"`
548+
ExpandedDirectory string `json:"expanded_directory"`
549+
Subsystem codersdk.AgentSubsystem `json:"subsystem"`
556550
}
557551

558552
func (c *Client) PostStartup(ctx context.Context, req PostStartupRequest) error {

‎codersdk/workspaceagents.go

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -130,9 +130,10 @@ type WorkspaceAgent struct {
130130
// LoginBeforeReady if true, the agent will delay logins until it is ready (e.g. executing startup script has ended).
131131
LoginBeforeReady bool `json:"login_before_ready"`
132132
// 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.
133-
StartupScriptTimeoutSeconds int32 `json:"startup_script_timeout_seconds"`
134-
ShutdownScript string `json:"shutdown_script,omitempty"`
135-
ShutdownScriptTimeoutSeconds int32 `json:"shutdown_script_timeout_seconds"`
133+
StartupScriptTimeoutSeconds int32 `json:"startup_script_timeout_seconds"`
134+
ShutdownScript string `json:"shutdown_script,omitempty"`
135+
ShutdownScriptTimeoutSeconds int32 `json:"shutdown_script_timeout_seconds"`
136+
Subsystem AgentSubsystem `json:"subsytem"`
136137
}
137138

138139
type DERPRegion struct {
@@ -553,3 +554,9 @@ type WorkspaceAgentStartupLog struct {
553554
Output string `json:"output"`
554555
Level LogLevel `json:"level"`
555556
}
557+
558+
type AgentSubsystem string
559+
560+
const (
561+
AgentSubsystemEnvbox AgentSubsystem = "envbox"
562+
)

0 commit comments

Comments
 (0)