diff --git a/enterprise/coderd/x/chatd/chatd_test.go b/enterprise/coderd/x/chatd/chatd_test.go index b959d7002f..a59a61780b 100644 --- a/enterprise/coderd/x/chatd/chatd_test.go +++ b/enterprise/coderd/x/chatd/chatd_test.go @@ -1471,6 +1471,14 @@ func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) { require.NoError(t, worker.Close()) }) + // Use a mock clock for the subscriber so the relay drain + // timer never fires until we explicitly advance it. This + // removes the nondeterministic 200ms race between the drain + // timer and the multi-hop snapshot forwarding pipeline. + subscriberClock := quartz.NewMock(t) + trapDrain := subscriberClock.Trap().NewTimer("drain") + defer trapDrain.Close() + // Subscriber dials through to the worker. On cancel the relay // drain fires well inside the worker's 5s grace, exercising the // cleanupStreamIfIdle early-return path. @@ -1490,7 +1498,7 @@ func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) { return nil, nil, nil, xerrors.New("worker subscribe failed") } return snapshot, relayEvents, cancel, nil - }, nil) + }, subscriberClock) ctx := testutil.Context(t, testutil.WaitLong) user, org, model := seedChatDependencies(ctx, t, db) @@ -1514,9 +1522,13 @@ func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) { // the assistant message and at least one message_part so we know // processChat's defer has flipped buffering=false and populated // bufferRetainedAt before the subscriber detaches. + // + // Each Eventually gets its own context so one slow assertion + // cannot starve subsequent ones of their deadline. var committedAssistantMsgs int var messagePartsSeen int - testutil.Eventually(ctx, t, func(context.Context) bool { + evCtx1 := testutil.Context(t, testutil.WaitLong) + testutil.Eventually(evCtx1, t, func(context.Context) bool { select { case event := <-events: switch event.Type { @@ -1533,7 +1545,15 @@ func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) { } }, testutil.IntervalFast) - testutil.Eventually(ctx, t, func(ctx context.Context) bool { + // Wait for the drain timer to be armed by the relay manager, + // then release and advance the subscriber clock past the drain + // timeout. This deterministically fires the drain without + // relying on wall-clock timing. + trapDrain.MustWait(ctx).MustRelease(ctx) + subscriberClock.Advance(200 * time.Millisecond).MustWait(ctx) + + evCtx2 := testutil.Context(t, testutil.WaitLong) + testutil.Eventually(evCtx2, t, func(ctx context.Context) bool { fromDB, dbErr := db.GetChatByID(ctx, chat.ID) if dbErr != nil { return false @@ -1550,7 +1570,8 @@ func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) { // worker observes the teardown. The retry itself re-enters // cleanupStreamIfIdle via its own cancel defer but still // early-returns because grace is still open. - testutil.Eventually(ctx, t, func(ctx context.Context) bool { + evCtx3 := testutil.Context(t, testutil.WaitLong) + testutil.Eventually(evCtx3, t, func(ctx context.Context) bool { snap, _, snapCancel, ok := worker.Subscribe(ctx, chat.ID, nil, math.MaxInt64) if !ok { return false