May 19, 2025 · May 19, 2025 · May 19, 2025 · May 19, 2025
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go

 db, ps := dbtestutil.NewDB(t)
 pubsubSpy := pubsubReinitSpy{
 Pubsub:     ps,
 subscribed: make(chan string),
 Pubsub:ps,
 triedToSubscribe: make(chan string),
 }
 client := coderdtest.New(t, &coderdtest.Options{
 Database: db,
 OwnerID:        user.UserID,
 }).WithAgent().Do()

 pubsubSpy.Mutex.Lock()
 pubsubSpy.Lock()
 pubsubSpy.expectedEvent = agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
 pubsubSpy.Mutex.Unlock()
 pubsubSpy.Unlock()

 agentCtx := testutil.Context(t, testutil.WaitShort)
 agentClient := agentsdk.New(client.URL)

 // 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
 testutil.TryReceive(ctx, t, pubsubSpy.triedToSubscribe)

 // Now that we're subscribed, publish the event
 err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
 type pubsubReinitSpy struct {
 pubsub.Pubsub
 sync.Mutex
 subscribed chan string
 expectedEvent string
 triedToSubscribe chan string
 expectedEventstring
 }

 func (p *pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
 cancel, err = p.Pubsub.Subscribe(event, listener)
 p.Lock()
 if p.expectedEvent != "" && event == p.expectedEvent {
 close(p.subscribed)
 close(p.triedToSubscribe)
 }
 p.Unlock()
 returnp.Pubsub.Subscribe(event, listener)
 returncancel, err
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -2650,8 +2650,8 @@ func TestReinit(t *testing.T) {

		db, ps := dbtestutil.NewDB(t)
		pubsubSpy := pubsubReinitSpy{
		Pubsub: ps,
		subscribed: make(chan string),
		Pubsub:ps,
		triedToSubscribe: make(chan string),
		}
		client := coderdtest.New(t, &coderdtest.Options{
		Database: db,
Expand All		@@ -2664,9 +2664,9 @@ func TestReinit(t *testing.T) {
		OwnerID: user.UserID,
		}).WithAgent().Do()

		pubsubSpy.Mutex.Lock()
		pubsubSpy.Lock()
		pubsubSpy.expectedEvent = agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
		pubsubSpy.Mutex.Unlock()
		pubsubSpy.Unlock()

		agentCtx := testutil.Context(t, testutil.WaitShort)
		agentClient := agentsdk.New(client.URL)
Expand All		@@ -2681,7 +2681,7 @@ func TestReinit(t *testing.T) {

		// 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
		testutil.TryReceive(ctx, t, pubsubSpy.triedToSubscribe)

		// Now that we're subscribed, publish the event
		err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
Expand All		@@ -2699,15 +2699,16 @@ func TestReinit(t *testing.T) {
		type pubsubReinitSpy struct {
		pubsub.Pubsub
		sync.Mutex
		subscribed chan string
		expectedEvent string
		triedToSubscribe chan string
		expectedEventstring
		}

		func (p *pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
		cancel, err = p.Pubsub.Subscribe(event, listener)
		p.Lock()
Comment on lines +2707 to 2708 Copy link Member johnstcnMay 19, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. Would it not make more sense to move the subscribe call after we lock the spy's mutex? Copy link Contributor spikecurtisMay 19, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. I don't think so; the mutex is protecting the`expectedEvent` and`triedToSubscribe`, not the underlying pubsub.
		if p.expectedEvent != "" && event == p.expectedEvent {
		close(p.subscribed)
		close(p.triedToSubscribe)
		}
		p.Unlock()
		returnp.Pubsub.Subscribe(event, listener)
		returncancel, err
		}