opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(enterprise/coderd/x/chatd): deflake TestSubscribeRelayDrainWithinGraceLeavesBufferRetained (#24549)

Browse Source 
Fixes the flake reported in
https://github.com/coder/internal/issues/1474.

- Use a `quartz.NewMock` clock for the subscriber with a drain timer
trap, so the 200ms relay drain fires only when explicitly advanced —
fully deterministic, no wall-clock race
- Give each `testutil.Eventually` call its own context so one slow
assertion cannot starve subsequent ones of their shared 25s deadline

> 🤖
This commit is contained in:
Cian Johnston
2026-04-21 11:20:32 +00:00
committed by GitHub
parent cb67e71835
commit f56adf5731
1 changed files with 25 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/x/chatd/chatd_test.go
+25 -4
View File
@@ -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