Make TestReinitializeAgent more robust

Minor code quality enhancements
coder · SasSwart · May 14, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025
commit 5363dcc078c6dd2c6e83d58c42279a8fde9d9997
diff --git a/cli/agent.go b/cli/agent.go
@@ -377,7 +377,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 					mustExit = true
 				case event := <-reinitEvents:
 					logger.Info(ctx, "agent received instruction to reinitialize",
-						slog.F("user_id", event.UserID), slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
+						slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
 				}
 
 				lastErr = agnt.Close()

diff --git a/coderd/prebuilds/claim.go b/coderd/prebuilds/claim.go
@@ -22,7 +22,7 @@ type PubsubWorkspaceClaimPublisher struct {
 
 func (p PubsubWorkspaceClaimPublisher) PublishWorkspaceClaim(claim agentsdk.ReinitializationEvent) error {
 	channel := agentsdk.PrebuildClaimedChannel(claim.WorkspaceID)
-	if err := p.ps.Publish(channel, []byte(claim.UserID.String())); err != nil {
+	if err := p.ps.Publish(channel, []byte(claim.Reason)); err != nil {
 		return xerrors.Errorf("failed to trigger prebuilt workspace agent reinitialization: %w", err)
 	}
 	return nil
@@ -48,16 +48,10 @@ func (p PubsubWorkspaceClaimListener) ListenForWorkspaceClaims(ctx context.Conte
 	default:
 	}
 
-	cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, id []byte) {
-		claimantID, err := uuid.ParseBytes(id)
-		if err != nil {
-			p.logger.Error(ctx, "invalid prebuild claimed channel payload", slog.F("input", string(id)))
-			return
-		}
+	cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, reason []byte) {
 		claim := agentsdk.ReinitializationEvent{
-			UserID:      claimantID,
 			WorkspaceID: workspaceID,
-			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+			Reason:      agentsdk.ReinitializationReason(reason),
 		}
 
 		select {

diff --git a/coderd/prebuilds/claim_test.go b/coderd/prebuilds/claim_test.go
@@ -25,7 +25,6 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
 		logger := testutil.Logger(t)
 		ps := pubsub.NewInMemory()
 		workspaceID := uuid.New()
-		userID := uuid.New()
 		reinitEvents := make(chan agentsdk.ReinitializationEvent, 1)
 		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
 		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, logger)
@@ -35,15 +34,15 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
 		defer cancel()
 
 		claim := agentsdk.ReinitializationEvent{
-			UserID:      userID,
 			WorkspaceID: workspaceID,
 			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 		}
 		err = publisher.PublishWorkspaceClaim(claim)
 		require.NoError(t, err)
 
-		gotUserID := testutil.RequireReceive(testutil.Context(t, testutil.WaitShort), t, reinitEvents)
-		require.Equal(t, userID, gotUserID.UserID)
+		gotEvent := testutil.RequireReceive(ctx, t, reinitEvents)
+		require.Equal(t, workspaceID, gotEvent.WorkspaceID)
+		require.Equal(t, claim.Reason, gotEvent.Reason)
 	})
 
 	t.Run("fail to publish claim", func(t *testing.T) {
@@ -53,7 +52,6 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
 
 		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
 		claim := agentsdk.ReinitializationEvent{
-			UserID:      uuid.New(),
 			WorkspaceID: uuid.New(),
 			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 		}
@@ -74,26 +72,21 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
 		claims := make(chan agentsdk.ReinitializationEvent, 1) // Buffer to avoid messing with goroutines in the rest of the test
 
 		workspaceID := uuid.New()
-		userID := uuid.New()
 		cancelFunc, err := listener.ListenForWorkspaceClaims(context.Background(), workspaceID, claims)
 		require.NoError(t, err)
 		defer cancelFunc()
 
 		// Publish a claim
 		channel := agentsdk.PrebuildClaimedChannel(workspaceID)
-		err = ps.Publish(channel, []byte(userID.String()))
+		reason := agentsdk.ReinitializeReasonPrebuildClaimed
+		err = ps.Publish(channel, []byte(reason))
 		require.NoError(t, err)
 
 		// Verify we receive the claim
-		select {
-		case claim, ok := <-claims:
-			require.True(t, ok, "received on a closed channel")
-			require.Equal(t, userID, claim.UserID)
-			require.Equal(t, workspaceID, claim.WorkspaceID)
-			require.Equal(t, agentsdk.ReinitializeReasonPrebuildClaimed, claim.Reason)
-		case <-time.After(testutil.WaitSuperLong): // TODO: revert to waitshort
-			t.Fatal("timeout waiting for claim")
-		}
+		ctx := testutil.Context(t, testutil.WaitShort)
+		claim := testutil.RequireReceive(ctx, t, claims)
+		require.Equal(t, workspaceID, claim.WorkspaceID)
+		require.Equal(t, reason, claim.Reason)
 	})
 
 	t.Run("ignores claim events for other workspaces", func(t *testing.T) {
@@ -111,7 +104,7 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
 
 		// Publish a claim for a different workspace
 		channel := agentsdk.PrebuildClaimedChannel(otherWorkspaceID)
-		err = ps.Publish(channel, []byte(uuid.New().String()))
+		err = ps.Publish(channel, []byte(agentsdk.ReinitializeReasonPrebuildClaimed))
 		require.NoError(t, err)
 
 		// Verify we don't receive the claim

diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -169,8 +169,12 @@ func TestAcquireJob(t *testing.T) {
 			_, err = tc.acquire(ctx, srv)
 			require.ErrorContains(t, err, "sql: no rows in result set")
 		})
-		for _, prebuiltWorkspace := range []bool{false, true} {
-			prebuiltWorkspace := prebuiltWorkspace
+		for _, prebuiltWorkspaceBuildStage := range []sdkproto.PrebuiltWorkspaceBuildStage{
+			sdkproto.PrebuiltWorkspaceBuildStage_NONE,
+			sdkproto.PrebuiltWorkspaceBuildStage_CREATE,
+			sdkproto.PrebuiltWorkspaceBuildStage_CLAIM,
+		} {
+			prebuiltWorkspaceBuildStage := prebuiltWorkspaceBuildStage
 			t.Run(tc.name+"_WorkspaceBuildJob", func(t *testing.T) {
 				t.Parallel()
 				// Set the max session token lifetime so we can assert we
@@ -294,31 +298,43 @@ func TestAcquireJob(t *testing.T) {
 					OwnerID:        user.ID,
 					OrganizationID: pd.OrganizationID,
 				})
-				build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+				build := database.WorkspaceBuild{
 					WorkspaceID:       workspace.ID,
 					BuildNumber:       1,
 					JobID:             uuid.New(),
 					TemplateVersionID: version.ID,
 					Transition:        database.WorkspaceTransitionStart,
 					Reason:            database.BuildReasonInitiator,
-				})
-				var buildState sdkproto.PrebuiltWorkspaceBuildStage
-				if prebuiltWorkspace {
-					buildState = sdkproto.PrebuiltWorkspaceBuildStage_CREATE
 				}
-				_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
-					ID:             build.ID,
+				input := provisionerdserver.WorkspaceProvisionJob{
+					WorkspaceBuildID: build.ID,
+				}
+				dbJob := database.ProvisionerJob{
+					ID:             build.JobID,
 					OrganizationID: pd.OrganizationID,
 					InitiatorID:    user.ID,
 					Provisioner:    database.ProvisionerTypeEcho,
 					StorageMethod:  database.ProvisionerStorageMethodFile,
 					FileID:         file.ID,
 					Type:           database.ProvisionerJobTypeWorkspaceBuild,
-					Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
-						WorkspaceBuildID:            build.ID,
-						PrebuiltWorkspaceBuildStage: buildState,
-					})),
-				})
+					Input:          must(json.Marshal(input)),
+				}
+				if prebuiltWorkspaceBuildStage == sdkproto.PrebuiltWorkspaceBuildStage_CLAIM {
+					dbgen.WorkspaceBuild(t, db, build)
+					prebuildJob := dbgen.ProvisionerJob(t, db, ps, dbJob)
+					resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+						JobID: prebuildJob.ID,
+					})
+					dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+						ResourceID: resource.ID,
+						AuthToken:  uuid.New(),
+					})
+					build.BuildNumber = 2
+					input.PrebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CLAIM
+					dbJob.Input = must(json.Marshal(input))
+				}
+				build = dbgen.WorkspaceBuild(t, db, build)
+				dbJob = dbgen.ProvisionerJob(t, db, ps, dbJob)
 
 				startPublished := make(chan struct{})
 				var closed bool
@@ -352,7 +368,7 @@ func TestAcquireJob(t *testing.T) {
 
 				<-startPublished
 
-				got, err := json.Marshal(job.Type)
+				got, err := json.Marshal(dbJob.Type)
 				require.NoError(t, err)
 
 				// Validate that a session token is generated during the job.
@@ -385,9 +401,7 @@ func TestAcquireJob(t *testing.T) {
 					WorkspaceBuildId:              build.ID.String(),
 					WorkspaceOwnerLoginType:       string(user.LoginType),
 					WorkspaceOwnerRbacRoles:       []*sdkproto.Role{{Name: rbac.RoleOrgMember(), OrgId: pd.OrganizationID.String()}, {Name: "member", OrgId: ""}, {Name: rbac.RoleOrgAuditor(), OrgId: pd.OrganizationID.String()}},
-				}
-				if prebuiltWorkspace {
-					wantedMetadata.PrebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CREATE
+					PrebuiltWorkspaceBuildStage:   prebuiltWorkspaceBuildStage,
 				}
 
 				slices.SortFunc(wantedMetadata.WorkspaceOwnerRbacRoles, func(a, b *sdkproto.Role) int {

diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -2647,62 +2647,60 @@ func TestAgentConnectionInfo(t *testing.T) {
 func TestReinit(t *testing.T) {
 	t.Parallel()
 
-	t.Run("claimed workspaces are reinitialized", func(t *testing.T) {
-		t.Parallel()
-
-		db, ps := dbtestutil.NewDB(t)
-		pubsubSpy := pubsubReinitSpy{
-			Pubsub:        ps,
-			subscriptions: make(chan string),
-		}
-		client := coderdtest.New(t, &coderdtest.Options{
-			Database: db,
-			Pubsub:   pubsubSpy,
-		})
-		user := coderdtest.CreateFirstUser(t, client)
-
-		r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
-			OrganizationID: user.OrganizationID,
-			OwnerID:        user.UserID,
-		}).WithAgent().Do()
+	db, ps := dbtestutil.NewDB(t)
+	pubsubSpy := pubsubReinitSpy{
+		Pubsub:     ps,
+		subscribed: make(chan string),
+	}
+	client := coderdtest.New(t, &coderdtest.Options{
+		Database: db,
+		Pubsub:   &pubsubSpy,
+	})
+	user := coderdtest.CreateFirstUser(t, client)
 
-		agentCtx := testutil.Context(t, testutil.WaitShort)
-		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(r.AgentToken)
+	r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+		OrganizationID: user.OrganizationID,
+		OwnerID:        user.UserID,
+	}).WithAgent().Do()
+	pubsubSpy.expectedEvent = agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
 
-		agentReinitializedCh := make(chan *agentsdk.ReinitializationEvent)
-		go func() {
-			reinitEvent, err := agentClient.WaitForReinit(agentCtx)
-			assert.NoError(t, err)
-			agentReinitializedCh <- reinitEvent
-		}()
+	agentCtx := testutil.Context(t, testutil.WaitShort)
+	agentClient := agentsdk.New(client.URL)
+	agentClient.SetSessionToken(r.AgentToken)
 
-		// We need to subscribe before we publish, lest we miss the event
-		for subscription := range pubsubSpy.subscriptions {
-			if subscription == agentsdk.PrebuildClaimedChannel(r.Workspace.ID) {
-				break
-			}
-		}
+	agentReinitializedCh := make(chan *agentsdk.ReinitializationEvent)
+	go func() {
+		reinitEvent, err := agentClient.WaitForReinit(agentCtx)
+		assert.NoError(t, err)
+		agentReinitializedCh <- reinitEvent
+	}()
 
-		err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
-			WorkspaceID: r.Workspace.ID,
-			UserID:      user.UserID,
-		})
-		require.NoError(t, err)
+	// We need to subscribe before we publish, lest we miss the event
+	ctx := testutil.Context(t, testutil.WaitShort)
+	testutil.TryReceive(ctx, t, pubsubSpy.subscribed) // Wait for the appropriate subscription
 
-		ctx := testutil.Context(t, testutil.WaitShort)
-		reinitEvent := testutil.TryReceive(ctx, t, agentReinitializedCh)
-		require.NotNil(t, reinitEvent)
-		require.Equal(t, r.Workspace.ID, reinitEvent.WorkspaceID)
+	// Now that we're subscribed, publish the event
+	err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
+		WorkspaceID: r.Workspace.ID,
+		Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 	})
+	require.NoError(t, err)
+
+	ctx = testutil.Context(t, testutil.WaitShort)
+	reinitEvent := testutil.TryReceive(ctx, t, agentReinitializedCh)
+	require.NotNil(t, reinitEvent)
+	require.Equal(t, r.Workspace.ID, reinitEvent.WorkspaceID)
 }
 
 type pubsubReinitSpy struct {
 	pubsub.Pubsub
-	subscriptions chan string
+	subscribed    chan string
+	expectedEvent string
 }
 
-func (p pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
-	p.subscriptions <- event
+func (p *pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
+	if p.expectedEvent != "" && event == p.expectedEvent {
+		close(p.subscribed)
+	}
 	return p.Pubsub.Subscribe(event, listener)
 }
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -698,7 +698,6 @@ const (
 
 type ReinitializationEvent struct {
 	WorkspaceID uuid.UUID
-	UserID      uuid.UUID
 	Reason      ReinitializationReason `json:"reason"`
 }
 
@@ -806,8 +805,9 @@ func (s *SSEAgentReinitTransmitter) Transmit(ctx context.Context, reinitEvents <
 		return xerrors.Errorf("failed to create sse transmitter: %w", err)
 	}
 
-	// Prevent handler from returning until the sender is closed.
 	defer func() {
+		// Block returning until the ServerSentEventSender is closed
+		// to avoid a race condition where we might write or flush to rw after the handler returns.
 		<-sseSenderClosed
 	}()
 

diff --git a/codersdk/agentsdk/agentsdk_test.go b/codersdk/agentsdk/agentsdk_test.go
@@ -22,7 +22,6 @@ func TestStreamAgentReinitEvents(t *testing.T) {
 		t.Parallel()
 
 		eventToSend := agentsdk.ReinitializationEvent{
-			UserID:      uuid.New(),
 			WorkspaceID: uuid.New(),
 			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 		}
@@ -56,7 +55,6 @@ func TestStreamAgentReinitEvents(t *testing.T) {
 		t.Parallel()
 
 		eventToSend := agentsdk.ReinitializationEvent{
-			UserID:      uuid.New(),
 			WorkspaceID: uuid.New(),
 			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 		}
@@ -92,7 +90,6 @@ func TestStreamAgentReinitEvents(t *testing.T) {
 		t.Parallel()
 
 		eventToSend := agentsdk.ReinitializationEvent{
-			UserID:      uuid.New(),
 			WorkspaceID: uuid.New(),
 			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 		}