provider_stream_test.go

package ai

import (
	"bytes"
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestProviders_SendStream_ParsesChunks(t *testing.T) {
	tests := []struct {
		name       string
		makeProv   func(baseURL string) Provider
		path       string
		writeBody  func(w http.ResponseWriter)
		wantFull   string
		wantTokens func(full string) (prompt, completion, total int)
	}{
		{
			name: "openai",
			makeProv: func(baseURL string) Provider {
				return NewOpenAIProvider(ProviderConfig{APIKey: "k", BaseURL: baseURL})
			},
			path: "/chat/completions",
			writeBody: func(w http.ResponseWriter) {
				_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"}}]}\n"))
				_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n"))
				_, _ = w.Write([]byte("data: [DONE]\n"))
			},
			wantFull: "Hello",
			wantTokens: func(full string) (int, int, int) {
				// streaming estimates completion tokens as len/4
				ct := len(full) / 4
				return 0, ct, ct
			},
		},
		{
			name: "openrouter",
			makeProv: func(baseURL string) Provider {
				return NewOpenRouterProvider(ProviderConfig{APIKey: "k", BaseURL: baseURL})
			},
			path: "/chat/completions",
			writeBody: func(w http.ResponseWriter) {
				_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"A\"}}]}\n"))
				_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"B\"}}]}\n"))
				_, _ = w.Write([]byte("data: [DONE]\n"))
			},
			wantFull: "AB",
			wantTokens: func(full string) (int, int, int) {
				ct := len(full) / 4
				return 0, ct, ct
			},
		},
		{
			name: "anthropic",
			makeProv: func(baseURL string) Provider {
				return NewAnthropicProvider(ProviderConfig{APIKey: "k", BaseURL: baseURL})
			},
			path: "/messages",
			writeBody: func(w http.ResponseWriter) {
				_, _ = w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Hi\"}}\n"))
				_, _ = w.Write([]byte("data: {\"type\":\"message_stop\"}\n"))
			},
			wantFull: "Hi",
			wantTokens: func(full string) (int, int, int) {
				ct := len(full) / 4
				return 0, ct, ct
			},
		},
		{
			name: "google",
			makeProv: func(baseURL string) Provider {
				return NewGoogleProvider(ProviderConfig{APIKey: "k", BaseURL: baseURL})
			},
			// google uses full URL with query params; we'll match by prefix.
			path: "/models/gemini-3-flash-preview:streamGenerateContent",
			writeBody: func(w http.ResponseWriter) {
				_, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Yo\"}]}}]}\n"))
				_, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"!\"}]}}]}\n"))
			},
			wantFull: "Yo!",
			wantTokens: func(full string) (int, int, int) {
				ct := len(full) / 4
				return 0, ct, ct
			},
		},
		{
			name: "ollama",
			makeProv: func(baseURL string) Provider {
				return NewOllamaProvider(ProviderConfig{BaseURL: baseURL})
			},
			path: "/api/chat",
			writeBody: func(w http.ResponseWriter) {
				_, _ = w.Write([]byte("{\"model\":\"m\",\"message\":{\"role\":\"assistant\",\"content\":\"He\"},\"done\":false}\n"))
				_, _ = w.Write([]byte("{\"model\":\"m\",\"message\":{\"role\":\"assistant\",\"content\":\"y\"},\"done\":true,\"prompt_eval_count\":2,\"eval_count\":3}\n"))
			},
			wantFull: "Hey",
			wantTokens: func(full string) (int, int, int) {
				return 2, 3, 5
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cleanup := withTestGlobals(t)
			defer cleanup()

			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				switch tt.name {
				case "google":
					if !strings.HasPrefix(r.URL.Path, tt.path) {
						t.Fatalf("expected path prefix %q, got %q", tt.path, r.URL.Path)
					}
				default:
					if r.URL.Path != tt.path {
						t.Fatalf("expected path %q, got %q", tt.path, r.URL.Path)
					}
				}
				w.Header().Set("Content-Type", "text/event-stream")
				tt.writeBody(w)
			}))
			defer srv.Close()

			p := tt.makeProv(srv.URL)
			var chunks []string
			resp, err := p.SendStream(context.Background(), &ProviderRequest{
				Model:    string(ModelGemini3Flash),
				Messages: []Message{{Role: "user", Content: "hi"}},
				Stream:   true,
			}, func(chunk string) {
				chunks = append(chunks, chunk)
			})
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			if resp.Content != tt.wantFull {
				t.Fatalf("expected full content %q, got %q", tt.wantFull, resp.Content)
			}

			joined := strings.Join(chunks, "")
			if joined != tt.wantFull {
				t.Fatalf("expected streamed chunks %q, got %q", tt.wantFull, joined)
			}

			// Ensure callback ordering is preserved.
			if bytes.Join(func() [][]byte {
				out := make([][]byte, 0, len(chunks))
				for _, c := range chunks {
					out = append(out, []byte(c))
				}
				return out
			}(), nil); false {
				// no-op: keeps imports stable if you tweak chunk checks later
			}

			pt, ct, ttok := tt.wantTokens(tt.wantFull)
			if resp.PromptTokens != pt || resp.CompletionTokens != ct || resp.TotalTokens != ttok {
				t.Fatalf("unexpected tokens: prompt=%d completion=%d total=%d", resp.PromptTokens, resp.CompletionTokens, resp.TotalTokens)
			}
		})
	}
}