test: fix a race in TestReinit (#17902)

closes https://github.com/coder/internal/issues/632

`pubsubReinitSpy` used to signal that a subscription had happened before
it actually had.
This created a slight opportunity for the main goroutine to publish
before the actual subscription was listening. The published event was
then dropped, leading to a failed test.
This commit is contained in:
Sas Swart
2025-05-19 11:37:54 +02:00
committed by GitHub
parent 1a41608035
commit c775ea8411

View File

@ -2650,8 +2650,8 @@ func TestReinit(t *testing.T) {
db, ps := dbtestutil.NewDB(t) db, ps := dbtestutil.NewDB(t)
pubsubSpy := pubsubReinitSpy{ pubsubSpy := pubsubReinitSpy{
Pubsub: ps, Pubsub: ps,
subscribed: make(chan string), triedToSubscribe: make(chan string),
} }
client := coderdtest.New(t, &coderdtest.Options{ client := coderdtest.New(t, &coderdtest.Options{
Database: db, Database: db,
@ -2664,9 +2664,9 @@ func TestReinit(t *testing.T) {
OwnerID: user.UserID, OwnerID: user.UserID,
}).WithAgent().Do() }).WithAgent().Do()
pubsubSpy.Mutex.Lock() pubsubSpy.Lock()
pubsubSpy.expectedEvent = agentsdk.PrebuildClaimedChannel(r.Workspace.ID) pubsubSpy.expectedEvent = agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
pubsubSpy.Mutex.Unlock() pubsubSpy.Unlock()
agentCtx := testutil.Context(t, testutil.WaitShort) agentCtx := testutil.Context(t, testutil.WaitShort)
agentClient := agentsdk.New(client.URL) agentClient := agentsdk.New(client.URL)
@ -2681,7 +2681,7 @@ func TestReinit(t *testing.T) {
// We need to subscribe before we publish, lest we miss the event // We need to subscribe before we publish, lest we miss the event
ctx := testutil.Context(t, testutil.WaitShort) 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 // Now that we're subscribed, publish the event
err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{ err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
@ -2699,15 +2699,16 @@ func TestReinit(t *testing.T) {
type pubsubReinitSpy struct { type pubsubReinitSpy struct {
pubsub.Pubsub pubsub.Pubsub
sync.Mutex sync.Mutex
subscribed chan string triedToSubscribe chan string
expectedEvent string expectedEvent string
} }
func (p *pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) { func (p *pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
cancel, err = p.Pubsub.Subscribe(event, listener)
p.Lock() p.Lock()
if p.expectedEvent != "" && event == p.expectedEvent { if p.expectedEvent != "" && event == p.expectedEvent {
close(p.subscribed) close(p.triedToSubscribe)
} }
p.Unlock() p.Unlock()
return p.Pubsub.Subscribe(event, listener) return cancel, err
} }