TestPromoteQueuedWhileRunning guards against the data-loss failure mode: promoting on a streaming chat must preserve partial assistant output by deferring the user-message insert to the worker's auto-promote.
(t *testing.T)
| 10520 | // partial assistant output by deferring the user-message insert |
| 10521 | // to the worker's auto-promote. |
| 10522 | func TestPromoteQueuedWhileRunning(t *testing.T) { |
| 10523 | t.Parallel() |
| 10524 | |
| 10525 | db, ps := dbtestutil.NewDB(t) |
| 10526 | ctx := testutil.Context(t, testutil.WaitLong) |
| 10527 | |
| 10528 | streamStarted := make(chan struct{}) |
| 10529 | streamCanceled := make(chan struct{}) |
| 10530 | var streamCallCount atomic.Int32 |
| 10531 | openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { |
| 10532 | if !req.Stream { |
| 10533 | return chattest.OpenAINonStreamingResponse("running-promote") |
| 10534 | } |
| 10535 | if streamCallCount.Add(1) > 1 { |
| 10536 | // Subsequent calls are the resumed run; let it settle. |
| 10537 | return chattest.OpenAIStreamingResponse( |
| 10538 | chattest.OpenAITextChunks("resumed after promotion")..., |
| 10539 | ) |
| 10540 | } |
| 10541 | chunks := make(chan chattest.OpenAIChunk, 1) |
| 10542 | go func() { |
| 10543 | defer close(chunks) |
| 10544 | chunks <- chattest.OpenAITextChunks("partial-running-output")[0] |
| 10545 | select { |
| 10546 | case <-streamStarted: |
| 10547 | default: |
| 10548 | close(streamStarted) |
| 10549 | } |
| 10550 | <-req.Context().Done() |
| 10551 | select { |
| 10552 | case <-streamCanceled: |
| 10553 | default: |
| 10554 | close(streamCanceled) |
| 10555 | } |
| 10556 | }() |
| 10557 | return chattest.OpenAIResponse{StreamingChunks: chunks} |
| 10558 | }) |
| 10559 | |
| 10560 | server := newActiveTestServer(t, db, ps) |
| 10561 | user, org, model := seedChatDependencies(t, db) |
| 10562 | setOpenAIProviderBaseURL(ctx, t, db, openAIURL) |
| 10563 | |
| 10564 | chat, err := server.CreateChat(ctx, chatd.CreateOptions{ |
| 10565 | OwnerID: user.ID, |
| 10566 | OrganizationID: org.ID, |
| 10567 | Title: "promote-while-running", |
| 10568 | ModelConfigID: model.ID, |
| 10569 | InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, |
| 10570 | }) |
| 10571 | require.NoError(t, err) |
| 10572 | |
| 10573 | testutil.Eventually(ctx, t, func(ctx context.Context) bool { |
| 10574 | fromDB, dbErr := db.GetChatByID(ctx, chat.ID) |
| 10575 | if dbErr != nil { |
| 10576 | return false |
| 10577 | } |
| 10578 | return fromDB.Status == database.ChatStatusRunning && fromDB.WorkerID.Valid |
| 10579 | }, testutil.IntervalFast) |
nothing calls this directly
no test coverage detected