mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(coderd/x/chatd): stabilize subagent pubsub completion test (#23944)
- stabilize `TestAwaitSubagentCompletion/CompletesViaPubsub` by waiting
for durable completion state before sending the synthetic pubsub wake
- add coverage for successful subagent completion with an empty report
> 🤖 Written by a Coder Agent. Reviewed by a human.
This commit is contained in:
@@ -1012,8 +1012,10 @@ func TestAwaitSubagentCompletion(t *testing.T) {
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
|
||||
|
||||
// Trap the fallback poll ticker to know when the
|
||||
// function has subscribed to pubsub and entered
|
||||
// its select loop.
|
||||
// function has entered the wait setup path. We still
|
||||
// need an explicit subscription handshake below because
|
||||
// the ticker can be created before SubscribeWithErr has
|
||||
// finished registering the listener.
|
||||
tickTrap := mClock.Trap().NewTicker("chatd", "subagent_poll")
|
||||
|
||||
type awaitResult struct {
|
||||
@@ -1029,19 +1031,47 @@ func TestAwaitSubagentCompletion(t *testing.T) {
|
||||
resultCh <- awaitResult{chat, report, err}
|
||||
}()
|
||||
|
||||
// Wait for the ticker to be created (confirms pubsub
|
||||
// subscription is set up and select loop entered).
|
||||
// Wait for the ticker to be created so the waiter has
|
||||
// entered its setup path, then subscribe our own probe on
|
||||
// the same channel. Because MemoryPubsub publishes only to
|
||||
// listeners already present at Publish time, waiting for
|
||||
// our probe to receive a message proves the waiter's
|
||||
// subscription is also registered before we assert on the
|
||||
// wake-up behavior.
|
||||
tickTrap.MustWait(ctx).MustRelease(ctx)
|
||||
tickTrap.Close()
|
||||
|
||||
// Transition child and publish. The pubsub notification
|
||||
// wakes the function without needing a clock advance.
|
||||
probeCh := make(chan struct{}, 1)
|
||||
cancelProbe, err := ps.SubscribeWithErr(
|
||||
coderdpubsub.ChatStreamNotifyChannel(child.ID),
|
||||
func(_ context.Context, _ []byte, _ error) {
|
||||
select {
|
||||
case probeCh <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
defer cancelProbe()
|
||||
|
||||
// Transition the child first, then publish once the
|
||||
// durable completion state is observable. Pubsub only
|
||||
// wakes the waiter; it does not guarantee the report is
|
||||
// visible in the same instant as the notification.
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
||||
insertAssistantMessage(ctx, t, db, child.ID, model.ID, "pubsub result")
|
||||
_ = ps.Publish(
|
||||
require.EventuallyWithT(t, func(c *assert.CollectT) {
|
||||
chat, report, done, err := server.checkSubagentCompletion(ctx, child.ID)
|
||||
require.NoError(c, err)
|
||||
assert.True(c, done)
|
||||
assert.Equal(c, child.ID, chat.ID)
|
||||
assert.Equal(c, "pubsub result", report)
|
||||
}, testutil.WaitMedium, testutil.IntervalFast)
|
||||
require.NoError(t, ps.Publish(
|
||||
coderdpubsub.ChatStreamNotifyChannel(child.ID),
|
||||
[]byte("done"),
|
||||
)
|
||||
))
|
||||
testutil.RequireReceive(ctx, t, probeCh)
|
||||
|
||||
result := testutil.RequireReceive(ctx, t, resultCh)
|
||||
require.NoError(t, result.err)
|
||||
@@ -1049,6 +1079,27 @@ func TestAwaitSubagentCompletion(t *testing.T) {
|
||||
assert.Equal(t, "pubsub result", result.report)
|
||||
})
|
||||
|
||||
t.Run("AlreadyWaitingNoReport", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
parent, child := createParentChildChats(ctx, t, server, user, model)
|
||||
|
||||
// signalWake from CreateChat may trigger immediate processing.
|
||||
// Wait for it to settle, then set the terminal state we need.
|
||||
// This case should return immediately, so use the shared
|
||||
// real-clock server instead of a mock clock.
|
||||
server.inflight.Wait()
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
||||
|
||||
gotChat, report, err := server.awaitSubagentCompletion(
|
||||
ctx, parent.ID, child.ID, 5*time.Second,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, child.ID, gotChat.ID)
|
||||
assert.Empty(t, report)
|
||||
})
|
||||
|
||||
t.Run("Timeout", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user