Skip to content

Commit 26dbc3a

Browse files
SasSwartSasSwart
committed
test agent reinitialization
1 parent c06c486 commit 26dbc3a

File tree

5 files changed

+112
-6
lines changed

5 files changed

+112
-6
lines changed

‎cli/agent.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -413,7 +413,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
413413
break
414414
}
415415

416-
logger.Info(ctx, "reinitializing...")
416+
logger.Info(ctx, "agent reinitializing")
417417
}
418418
return lastErr
419419
},

‎cli/agent_test.go

+52
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ import (
2121
"github.com/coder/coder/v2/coderd/coderdtest"
2222
"github.com/coder/coder/v2/coderd/database"
2323
"github.com/coder/coder/v2/coderd/database/dbfake"
24+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2425
"github.com/coder/coder/v2/codersdk"
26+
"github.com/coder/coder/v2/codersdk/agentsdk"
2527
"github.com/coder/coder/v2/codersdk/workspacesdk"
2628
"github.com/coder/coder/v2/provisionersdk/proto"
2729
"github.com/coder/coder/v2/testutil"
@@ -321,6 +323,56 @@ func TestWorkspaceAgent(t *testing.T) {
321323
})
322324
}
323325

326+
func TestAgent_Prebuild(t *testing.T) {
327+
t.Parallel()
328+
329+
db, pubsub := dbtestutil.NewDB(t)
330+
client := coderdtest.New(t, &coderdtest.Options{
331+
Database: db,
332+
Pubsub: pubsub,
333+
})
334+
user := coderdtest.CreateFirstUser(t, client)
335+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
336+
OrganizationID: user.OrganizationID,
337+
OwnerID: user.UserID,
338+
}).WithAgent(func(a []*proto.Agent) []*proto.Agent {
339+
a[0].Scripts = []*proto.Script{
340+
{
341+
DisplayName: "Prebuild Test Script",
342+
Script: "sleep 5", // Make reinitiazation take long enough to assert that it happened
343+
RunOnStart: true,
344+
},
345+
}
346+
return a
347+
}).Do()
348+
349+
// Spin up an agent
350+
logDir := t.TempDir()
351+
inv, _ := clitest.New(t,
352+
"agent",
353+
"--auth", "token",
354+
"--agent-token", r.AgentToken,
355+
"--agent-url", client.URL.String(),
356+
"--log-dir", logDir,
357+
)
358+
clitest.Start(t, inv)
359+
360+
// Check that the agent is in a happy steady state
361+
waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID)
362+
waiter.WaitFor(coderdtest.AgentReady)
363+
364+
// Trigger reinitialization
365+
channel := agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
366+
err := pubsub.Publish(channel, []byte(user.UserID.String()))
367+
require.NoError(t, err)
368+
369+
// Check that the agent reinitializes
370+
waiter.WaitFor(coderdtest.AgentNotReady)
371+
372+
// Check that reinitialization completed
373+
waiter.WaitFor(coderdtest.AgentReady)
374+
}
375+
324376
func matchAgentWithVersion(rs []codersdk.WorkspaceResource) bool {
325377
if len(rs) < 1 {
326378
return false

‎coderd/coderdtest/coderdtest.go

+53
Original file line numberDiff line numberDiff line change
@@ -1105,6 +1105,59 @@ func (w WorkspaceAgentWaiter) MatchResources(m func([]codersdk.WorkspaceResource
11051105
return w
11061106
}
11071107

1108+
type WaitForCriterium func(agent codersdk.WorkspaceAgent) bool
1109+
1110+
func AgentReady(agent codersdk.WorkspaceAgent) bool {
1111+
return agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady
1112+
}
1113+
1114+
func AgentNotReady(agent codersdk.WorkspaceAgent) bool {
1115+
return !AgentReady(agent)
1116+
}
1117+
1118+
func (w WorkspaceAgentWaiter) WaitFor(criteria ...WaitForCriterium) {
1119+
w.t.Helper()
1120+
1121+
agentNamesMap := make(map[string]struct{}, len(w.agentNames))
1122+
for _, name := range w.agentNames {
1123+
agentNamesMap[name] = struct{}{}
1124+
}
1125+
1126+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
1127+
defer cancel()
1128+
1129+
w.t.Logf("waiting for workspace agents (workspace %s)", w.workspaceID)
1130+
require.Eventually(w.t, func() bool {
1131+
var err error
1132+
workspace, err := w.client.Workspace(ctx, w.workspaceID)
1133+
if err != nil {
1134+
return false
1135+
}
1136+
if workspace.LatestBuild.Job.CompletedAt == nil {
1137+
return false
1138+
}
1139+
if workspace.LatestBuild.Job.CompletedAt.IsZero() {
1140+
return false
1141+
}
1142+
1143+
for _, resource := range workspace.LatestBuild.Resources {
1144+
for _, agent := range resource.Agents {
1145+
if len(w.agentNames) > 0 {
1146+
if _, ok := agentNamesMap[agent.Name]; !ok {
1147+
continue
1148+
}
1149+
}
1150+
for _, criterium := range criteria {
1151+
if !criterium(agent) {
1152+
return false
1153+
}
1154+
}
1155+
}
1156+
}
1157+
return true
1158+
}, testutil.WaitLong, testutil.IntervalMedium)
1159+
}
1160+
11081161
// Wait waits for the agent(s) to connect and fails the test if they do not within testutil.WaitLong
11091162
func (w WorkspaceAgentWaiter) Wait() []codersdk.WorkspaceResource {
11101163
w.t.Helper()

‎coderd/workspaces.go

+6-3
Original file line numberDiff line numberDiff line change
@@ -689,10 +689,13 @@ func createWorkspace(
689689
api.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
690690
slog.F("workspace_id", claimedWorkspace.ID), slog.Error(err))
691691
}
692-
agentTokensByAgentID = make(map[uuid.UUID]string, len(agents))
693-
for _, agent := range agents {
694-
agentTokensByAgentID[agent.ID] = agent.AuthToken.String()
692+
if len(agents) > 1 {
693+
return xerrors.Errorf("multiple agents are not yet supported in prebuilt workspaces")
695694
}
695+
// agentTokensByAgentID = make(map[uuid.UUID]string, len(agents))
696+
// for _, agent := range agents {
697+
// agentTokensByAgentID[agent.ID] = agent.AuthToken.String()
698+
// }
696699
}
697700

698701
// We have to refetch the workspace for the joined in fields.

‎go.sum

-2
Original file line numberDiff line numberDiff line change
@@ -921,8 +921,6 @@ github.com/coder/tailscale v1.1.1-0.20250422090654-5090e715905e h1:nope/SZfoLB9M
921921
github.com/coder/tailscale v1.1.1-0.20250422090654-5090e715905e/go.mod h1:1ggFFdHTRjPRu9Yc1yA7nVHBYB50w9Ce7VIXNqcW6Ko=
922922
github.com/coder/terraform-config-inspect v0.0.0-20250107175719-6d06d90c630e h1:JNLPDi2P73laR1oAclY6jWzAbucf70ASAvf5mh2cME0=
923923
github.com/coder/terraform-config-inspect v0.0.0-20250107175719-6d06d90c630e/go.mod h1:Gz/z9Hbn+4KSp8A2FBtNszfLSdT2Tn/uAKGuVqqWmDI=
924-
github.com/coder/terraform-provider-coder/v2 v2.4.0-pre1.0.20250417100258-c86bb5c3ddcd h1:FsIG6Fd0YOEK7D0Hl/CJywRA+Y6Gd5RQbSIa2L+/BmE=
925-
github.com/coder/terraform-provider-coder/v2 v2.4.0-pre1.0.20250417100258-c86bb5c3ddcd/go.mod h1:56/KdGYaA+VbwXJbTI8CA57XPfnuTxN8rjxbR34PbZw=
926924
github.com/coder/trivy v0.0.0-20250409153844-e6b004bc465a h1:yryP7e+IQUAArlycH4hQrjXQ64eRNbxsV5/wuVXHgME=
927925
github.com/coder/trivy v0.0.0-20250409153844-e6b004bc465a/go.mod h1:dDvq9axp3kZsT63gY2Znd1iwzfqDq3kXbQnccIrjRYY=
928926
github.com/coder/websocket v1.8.13 h1:f3QZdXy7uGVz+4uCJy2nTZyM0yTBj8yANEHhqlXZ9FE=

0 commit comments

Comments
 (0)