From 46472f1022d5318dec3e5e806b29153a849eede4 Mon Sep 17 00:00:00 2001 From: Lance Cooper Date: Sat, 22 Aug 2026 14:45:00 -0500 Subject: [PATCH 1/3] test: add stream helper infrastructure and tests for existing behavior --- internal/client/client_test.go | 218 +++++++++++++++++++++++++++++++++ 1 file changed, 218 insertions(+) diff --git a/internal/client/client_test.go b/internal/client/client_test.go index 416ca40..4f80302 100644 --- a/internal/client/client_test.go +++ b/internal/client/client_test.go @@ -1,7 +1,13 @@ package client import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" "testing" + "time" ) func TestSupportsVisionOverride(t *testing.T) { @@ -32,3 +38,215 @@ func TestSupportsVisionOverride(t *testing.T) { t.Errorf("expected SupportsVision() to be true when c.supportsVision is true") } } + +// Sample SSE chunk JSON strings shared across stream tests. +const ( + sampleChunkHello = `{"id":"c1","choices":[{"delta":{"content":"Hello"}}]}` + sampleChunkWorld = `{"id":"c1","choices":[{"delta":{"content":" world"}}]}` + sampleChunkStop = `{"id":"c1","choices":[{"delta":{"content":""},"finish_reason":"stop"}]}` +) + +// defaultRequest returns a minimal ChatCompletionRequest used by most stream tests. +func defaultRequest() ChatCompletionRequest { + return ChatCompletionRequest{ + Model: "test-model", + Messages: []ChatMessage{{Role: "user", Content: TextContent("hi")}}, + } +} + +// streamTest is a shared test fixture for ChatCompletionStream tests. +type streamTest struct { + server *httptest.Server + client *Client +} + +// newStreamTest creates a test server and client pair. The caller should call +// st.Handle() to set the response handler before making requests. +func newStreamTest(t *testing.T) *streamTest { + t.Helper() + st := &streamTest{ + server: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + })), + } + st.client = NewClient(Config{BaseURL: st.server.URL}) + return st +} + +// Handle replaces the server's response handler. +func (st *streamTest) Handle(handler http.HandlerFunc) { + st.server.Config.Handler = handler +} + +// Close shuts down the test server. Call via defer immediately after creation. +func (st *streamTest) Close() { + st.server.Close() +} + +// collectStream drains the output/error channels from ChatCompletionStream +// into a slice of ChatCompletionChunk and an optional error. +func collectStream(t *testing.T, ctx context.Context, c *Client, req ChatCompletionRequest) ([]ChatCompletionChunk, error) { + t.Helper() + outCh, errCh := c.ChatCompletionStream(ctx, req) + + var chunks []ChatCompletionChunk + var streamErr error + + // Collect in a goroutine to avoid deadlocks if channels don't close. + done := make(chan struct{}) + go func() { + defer close(done) + for chunk := range outCh { + chunks = append(chunks, chunk) + } + select { + case err, ok := <-errCh: + if ok && err != nil { + streamErr = err + } + default: + } + }() + + // Wait with a timeout to catch hangs. + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("stream collection timed out — channels were not closed") + } + + return chunks, streamErr +} + +// sseChunks builds an SSE-formatted response body from JSON strings. +func sseChunks(jsons ...string) string { + var b strings.Builder + for _, j := range jsons { + fmt.Fprintf(&b, "data: %s\n", j) + } + return b.String() +} + +// accumulateContent concatenates the delta content strings from all choices across chunks. +func accumulateContent(chunks []ChatCompletionChunk) string { + var b strings.Builder + for _, ch := range chunks { + if len(ch.Choices) > 0 { + b.WriteString(ch.Choices[0].Delta.Content.String()) + } + } + return b.String() +} + +func TestChatCompletionStream_Termination(t *testing.T) { + tests := []struct { + name string + body string + wantChunks int + wantErr bool + }{ + { + name: "DONE_sentinel", + body: sseChunks(sampleChunkHello, "[DONE]"), + wantChunks: 1, + wantErr: false, + }, + { + name: "connection_close_no_sentinel", + body: sseChunks(sampleChunkHello), + wantChunks: 1, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + st := newStreamTest(t) + defer st.Close() + + st.Handle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, tt.body) + })) + + chunks, err := collectStream(t, context.Background(), st.client, defaultRequest()) + + if got := len(chunks); got != tt.wantChunks { + t.Errorf("got %d chunks, want %d", got, tt.wantChunks) + } + if (err != nil) != tt.wantErr { + t.Errorf("error = %v, wantErr = %v", err, tt.wantErr) + } + }) + } +} + +func TestChatCompletionStream_ContentAccumulation(t *testing.T) { + st := newStreamTest(t) + defer st.Close() + + chunkPayloads := []string{ + sampleChunkHello, + sampleChunkWorld, + sampleChunkStop, + } + + st.Handle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, sseChunks(chunkPayloads...), "\ndata: [DONE]\n\n") + })) + + resultChunks, err := collectStream(t, context.Background(), st.client, defaultRequest()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Verify chunk count (3 data lines, none are [DONE]) + if got := len(resultChunks); got != 3 { + t.Errorf("got %d chunks, want 3", got) + } + + // Verify content accumulation + if got := accumulateContent(resultChunks); got != "Hello world" { + t.Errorf("accumulated content = %q, want %q", got, "Hello world") + } + + // Verify finish reason on last chunk + if resultChunks[2].Choices[0].FinishReason != "stop" { + t.Errorf("finish_reason = %q, want %q", resultChunks[2].Choices[0].FinishReason, "stop") + } +} + +func TestChatCompletionStream_Non200Status(t *testing.T) { + st := newStreamTest(t) + defer st.Close() + + st.Handle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + fmt.Fprint(w, `{"error":{"message":"internal error"}}`) + })) + + _, err := collectStream(t, context.Background(), st.client, defaultRequest()) + if err == nil { + t.Error("expected error for 500 response, got nil") + } +} + +func TestChatCompletionStream_InvalidJSON(t *testing.T) { + st := newStreamTest(t) + defer st.Close() + + st.Handle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, sseChunks("{invalid json}", `{"id":"c1","choices":[{"delta":{"content":"valid"}}]}`, "[DONE]")) + })) + + chunks, err := collectStream(t, context.Background(), st.client, defaultRequest()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Should have 1 valid chunk (the invalid JSON was skipped). + if got := len(chunks); got != 1 { + t.Errorf("got %d chunks, want 1 (invalid JSON should be skipped)", got) + } +} From accb6a3d6b998196081149ddc5dbe083ef0a7b61 Mon Sep 17 00:00:00 2001 From: Lance Cooper Date: Sat, 22 Aug 2026 14:45:53 -0500 Subject: [PATCH 2/3] fix: handle empty stream terminator and propagate scanner errors --- internal/client/client.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/internal/client/client.go b/internal/client/client.go index 540cd80..c7e1e56 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -167,10 +167,17 @@ func (c *Client) ChatCompletionStream(ctx context.Context, req ChatCompletionReq continue } data := strings.TrimPrefix(line, "data: ") + + // Handle [DONE] sentinel (OpenAI standard) if data == "[DONE]" { break } + // Handle empty data line — some servers signal end this way + if data == "" { + break + } + var chunk ChatCompletionChunk if err := json.Unmarshal([]byte(data), &chunk); err != nil { continue @@ -181,6 +188,15 @@ func (c *Client) ChatCompletionStream(ctx context.Context, req ChatCompletionReq return } } + + // If the loop exited because scanner.Scan() returned false (connection closed) + // or an empty data line, check for read errors and propagate them. + if err := scanner.Err(); err != nil { + select { + case errCh <- fmt.Errorf("stream interrupted: %w", err): + default: + } + } }() return out, errCh From 42731c7090c58b660e099ea036047de29e06cce1 Mon Sep 17 00:00:00 2001 From: Lance Cooper Date: Sat, 22 Aug 2026 14:47:34 -0500 Subject: [PATCH 3/3] test: add tests for empty stream terminator and error propagation --- internal/client/client_test.go | 89 ++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/internal/client/client_test.go b/internal/client/client_test.go index 4f80302..f463c40 100644 --- a/internal/client/client_test.go +++ b/internal/client/client_test.go @@ -157,6 +157,18 @@ func TestChatCompletionStream_Termination(t *testing.T) { wantChunks: 1, wantErr: false, }, + { + name: "empty_data_line", + body: sseChunks(sampleChunkHello, ""), + wantChunks: 1, + wantErr: false, + }, + { + name: "only_empty_data_line", + body: "data: \n", + wantChunks: 0, + wantErr: false, + }, } for _, tt := range tests { @@ -250,3 +262,80 @@ func TestChatCompletionStream_InvalidJSON(t *testing.T) { t.Errorf("got %d chunks, want 1 (invalid JSON should be skipped)", got) } } + +func TestChatCompletionStream_ContextCancellation(t *testing.T) { + st := newStreamTest(t) + defer st.Close() + + // Separate channel to block the server handler. Using ctx.Done() here + // would create a race — we need to distinguish "client exited due to + // context cancel" from "client exited because server closed connection." + cancelled := make(chan struct{}) + defer close(cancelled) + + st.Handle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // DiscoverBackend probes /props before the actual request. Return + // 404 for non-completions paths so those probes don't hang on <-cancelled. + if r.URL.Path != "/v1/chat/completions" { + w.WriteHeader(http.StatusNotFound) + return + } + + flusher := w.(http.Flusher) // httptest.Server always implements Flusher + + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, sseChunks(sampleChunkHello)) + flusher.Flush() + + // Block until test teardown closes the cancelled channel. + <-cancelled + })) + + ctx, cancel := context.WithCancel(context.Background()) + + outCh, errCh := st.client.ChatCompletionStream(ctx, defaultRequest()) + + // Collect the first chunk — should arrive promptly before cancellation. + select { + case chunk := <-outCh: + if got := chunk.Choices[0].Delta.Content.String(); got != "Hello" { + t.Errorf("first chunk content = %q, want %q", got, "Hello") + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for first chunk") + } + + // Cancel the context to trigger cleanup. + cancel() + + // Verify output channel closes after cancellation. + select { + case _, ok := <-outCh: + if ok { + t.Error("output channel should be closed after context cancellation") + } + case <-time.After(5 * time.Second): + t.Fatal("output channel not closed after context cancellation") + } + + // Context cancellation interrupts the HTTP read, causing scanner.Err() to + // send a "stream interrupted" error on errCh before the deferred close fires. + select { + case err := <-errCh: + if err == nil { + t.Error("expected error on errCh after context cancellation, got nil") + } + case <-time.After(5 * time.Second): + t.Fatal("error channel did not produce a value after context cancellation") + } + + // Verify errCh is now closed. + select { + case _, ok := <-errCh: + if ok { + t.Error("error channel should be closed, but produced another value") + } + case <-time.After(5 * time.Second): + t.Fatal("error channel not closed after context cancellation") + } +}