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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 52 additions & 0 deletions
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

Lines changed: 53 additions & 0 deletions
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

Lines changed: 6 additions & 3 deletions
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

Lines changed: 0 additions & 2 deletions
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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy