tapir/internal/adapters/llm/client_test.go

package llm

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

// mockServer returns an OpenAI-compatible endpoint that always replies with the
// given assistant content.
func mockServer(t *testing.T, response string) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/chat/completions" {
			t.Errorf("path = %q, want /chat/completions", r.URL.Path)
		}
		if ct := r.Header.Get("Content-Type"); ct != "application/json" {
			t.Errorf("Content-Type = %q, want application/json", ct)
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"choices": []map[string]any{
				{"message": map[string]any{"role": "assistant", "content": response}},
			},
		})
	}))
}

func TestClient_Complete(t *testing.T) {
	srv := mockServer(t, "hello world")
	defer srv.Close()

	c := New(srv.URL, "", "test-model", 10*time.Second)
	got, err := c.Complete(context.Background(), "you are helpful", "say hello")
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if got != "hello world" {
		t.Errorf("got %q, want %q", got, "hello world")
	}
}

// TestClient_SendsMaxTokens guards Tapir's ADR-004 change to the copied client:
// it MUST send a positive max_tokens, or thinking models return empty content.
func TestClient_SendsMaxTokens(t *testing.T) {
	var body chatRequest
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(&body)
		_ = json.NewEncoder(w).Encode(map[string]any{
			"choices": []map[string]any{{"message": map[string]any{"content": "ok"}}},
		})
	}))
	defer srv.Close()

	c := New(srv.URL, "", "test-model", 10*time.Second)
	if _, err := c.Complete(context.Background(), "sys", "user"); err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if body.MaxTokens <= 0 {
		t.Errorf("max_tokens = %d, want > 0 (thinking models return empty content without it)", body.MaxTokens)
	}
}

// TestClient_WithMaxTokens overrides the completion budget — the summarizer caps
// it small so prompt + max_tokens fits a small-context model's window (8k).
func TestClient_WithMaxTokens(t *testing.T) {
	var body chatRequest
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(&body)
		_ = json.NewEncoder(w).Encode(map[string]any{
			"choices": []map[string]any{{"message": map[string]any{"content": "ok"}}},
		})
	}))
	defer srv.Close()

	c := New(srv.URL, "", "test-model", 10*time.Second, WithMaxTokens(1500))
	if _, err := c.Complete(context.Background(), "sys", "user"); err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if body.MaxTokens != 1500 {
		t.Errorf("max_tokens = %d, want 1500", body.MaxTokens)
	}
}

// TestClient_UsageHookRecordsTokens: the usage hook fires with the model and the
// prompt/completion token counts parsed from the response usage block.
func TestClient_UsageHookRecordsTokens(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{
			"choices": []map[string]any{{"message": map[string]any{"content": "ok"}}},
			"usage":   map[string]any{"prompt_tokens": 123, "completion_tokens": 45},
		})
	}))
	defer srv.Close()

	var gotModel string
	var gotPrompt, gotCompletion int
	c := New(srv.URL, "", "test-model", 10*time.Second, WithUsageHook(func(model string, p, comp int) {
		gotModel, gotPrompt, gotCompletion = model, p, comp
	}))
	if _, err := c.Complete(context.Background(), "sys", "user"); err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if gotModel != "test-model" || gotPrompt != 123 || gotCompletion != 45 {
		t.Errorf("usage hook got (%q, %d, %d), want (test-model, 123, 45)", gotModel, gotPrompt, gotCompletion)
	}
}

func TestClient_ReturnsErrorOnNon200(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "overloaded", http.StatusServiceUnavailable)
	}))
	defer srv.Close()

	c := New(srv.URL, "", "test-model", 10*time.Second)
	if _, err := c.Complete(context.Background(), "sys", "user"); err == nil {
		t.Fatal("want error on non-200, got nil")
	}
}

func TestClient_SendsAuthHeader(t *testing.T) {
	var gotAuth string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"choices": []map[string]any{{"message": map[string]any{"content": "ok"}}},
		})
	}))
	defer srv.Close()

	c := New(srv.URL, "my-key", "test-model", 10*time.Second)
	if _, err := c.Complete(context.Background(), "sys", "user"); err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if gotAuth != "Bearer my-key" {
		t.Errorf("Authorization = %q, want %q", gotAuth, "Bearer my-key")
	}
}

func TestClient_Retries429(t *testing.T) {
	calls := 0
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		if calls == 1 {
			w.Header().Set("Retry-After", "0")
			w.WriteHeader(http.StatusTooManyRequests)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{
			"choices": []map[string]any{{"message": map[string]any{"content": "retried"}}},
		})
	}))
	defer srv.Close()

	c := New(srv.URL, "", "test-model", 10*time.Second)
	got, err := c.Complete(context.Background(), "sys", "user")
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if got != "retried" {
		t.Errorf("got %q, want %q", got, "retried")
	}
	if calls != 2 {
		t.Errorf("calls = %d, want 2", calls)
	}
}