package chat import ( "context" "encoding/json" "errors" "io" "net/http" "net/http/httptest" "strings" "testing" "time" ) // ollamaServer mocks /api/chat + /api/tags. Captures last request body. func ollamaServer(t *testing.T, status int, respBody string) (*httptest.Server, *string) { t.Helper() captured := "" srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/tags": w.WriteHeader(200) _, _ = w.Write([]byte(`{"models":[{"name":"qwen3.5:latest"}]}`)) return case "/api/chat": bs, _ := io.ReadAll(r.Body) captured = string(bs) w.WriteHeader(status) _, _ = w.Write([]byte(respBody)) return default: w.WriteHeader(404) } })) t.Cleanup(srv.Close) return srv, &captured } func TestOllama_ChatHappyPath(t *testing.T) { resp := `{ "model": "qwen3.5:latest", "message": {"content": "ok"}, "done": true, "prompt_eval_count": 7, "eval_count": 3 }` srv, captured := ollamaServer(t, 200, resp) o := NewOllama(srv.URL, 5*time.Second) out, err := o.Chat(context.Background(), Request{ Model: "qwen3.5:latest", Messages: []Message{{Role: "user", Content: "hi"}}, }) if err != nil { t.Fatalf("Chat: %v", err) } if out.Content != "ok" { t.Errorf("Content = %q, want ok", out.Content) } if out.InputTokens != 7 || out.OutputTokens != 3 { t.Errorf("tokens = (%d, %d), want (7, 3)", out.InputTokens, out.OutputTokens) } // Verify captured request shape: stream=false, options.temperature // surfaced. var sent map[string]any if err := json.Unmarshal([]byte(*captured), &sent); err != nil { t.Fatalf("parse captured: %v", err) } if sent["stream"] != false { t.Errorf("stream should be false, got %v", sent["stream"]) } } func TestOllama_StripsExplicitPrefix(t *testing.T) { srv, captured := ollamaServer(t, 200, `{"message":{"content":""},"done":true}`) o := NewOllama(srv.URL, 5*time.Second) _, err := o.Chat(context.Background(), Request{Model: "ollama/qwen3.5:latest"}) if err != nil { t.Fatalf("Chat: %v", err) } var sent map[string]any _ = json.Unmarshal([]byte(*captured), &sent) if sent["model"] != "qwen3.5:latest" { t.Errorf("upstream model = %v, want qwen3.5:latest (prefix stripped)", sent["model"]) } } func TestOllama_FormatJSON(t *testing.T) { srv, captured := ollamaServer(t, 200, `{"message":{"content":"{}"},"done":true}`) o := NewOllama(srv.URL, 5*time.Second) _, _ = o.Chat(context.Background(), Request{Model: "qwen3.5:latest", Format: "json"}) if !strings.Contains(*captured, `"format":"json"`) { t.Errorf("Format=json should set top-level format=json; captured=%s", *captured) } } func TestOllama_Available(t *testing.T) { srv, _ := ollamaServer(t, 200, "{}") o := NewOllama(srv.URL, 5*time.Second) if !o.Available() { t.Errorf("server is up; Available should be true") } } func TestOllama_UpstreamError(t *testing.T) { srv, _ := ollamaServer(t, 500, `{"error":"out of memory"}`) o := NewOllama(srv.URL, 5*time.Second) _, err := o.Chat(context.Background(), Request{Model: "qwen3.5:latest"}) if !errors.Is(err, ErrUpstream) { t.Errorf("500 should ErrUpstream; got %v", err) } } func TestOllamaCloud_FlattenMessages(t *testing.T) { system, prompt := flattenMessages([]Message{ {Role: "system", Content: "You are helpful."}, {Role: "user", Content: "hi"}, {Role: "assistant", Content: "hello!"}, {Role: "user", Content: "how are you?"}, }) if system != "You are helpful." { t.Errorf("system = %q, want 'You are helpful.'", system) } if !strings.Contains(prompt, "User: hi") || !strings.Contains(prompt, "Assistant: hello!") || !strings.Contains(prompt, "User: how are you?") { t.Errorf("prompt missing role tags: %q", prompt) } }