ollama-proxy/internal/proxy/handler_test.go

package proxy

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Atte149/ollama-proxy/internal/config"
	"github.com/Atte149/ollama-proxy/internal/log"
)

// mockUpstream is a configurable test upstream. It responds to /api/chat and
// /v1/chat/completions; for other paths it returns 200 with a short body.
type mockUpstream struct {
	status      int32 // current status to return
	body        string
	allowAuth   string // if non-empty, require this Bearer token
	chunks      []string
	requestAuth atomic.Int32 // number of requests seen with each auth value
}

func (m *mockUpstream) handler(w http.ResponseWriter, r *http.Request) {
	auth := r.Header.Get("Authorization")
	if m.allowAuth != "" && auth != "Bearer "+m.allowAuth {
		http.Error(w, "bad auth", http.StatusUnauthorized)
		return
	}
	if m.status >= 500 || m.status == 429 {
		http.Error(w, "rate limited", int(m.status))
		return
	}
	if len(m.chunks) > 0 {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Transfer-Encoding", "chunked")
		w.WriteHeader(200)
		flusher, _ := w.(http.Flusher)
		for _, c := range m.chunks {
			_, _ = io.WriteString(w, c)
			if flusher != nil {
				flusher.Flush()
			}
		}
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(int(m.status))
	_, _ = io.WriteString(w, m.body)
}

func newTestHandler(t *testing.T, upstreamURL string, accounts []config.Account, retries int) *Handler {
	t.Helper()
	// Apply the mock upstream URL to any account without an explicit BaseURL.
	for i := range accounts {
		if accounts[i].BaseURL == "" {
			accounts[i].BaseURL = upstreamURL
		}
	}
	b := NewBalancer(accounts, 50*time.Millisecond)
	return NewHandler(b, upstreamURL, retries, log.New("debug"))
}

func TestHandler_ProxiesVersion(t *testing.T) {
	up := &mockUpstream{status: 200, body: `{"version":"0.1.42"}`}
	srv := httptest.NewServer(http.HandlerFunc(up.handler))
	defer srv.Close()

	h := newTestHandler(t, srv.URL, mkAccts(1), 1)
	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/api/version", nil)
	h.ServeHTTP(rec, req)

	if rec.Code != 200 {
		t.Fatalf("status = %d, want 200", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "0.1.42") {
		t.Errorf("body = %q, want version", rec.Body.String())
	}
}

func TestHandler_FailoverOn429(t *testing.T) {
	// Two upstreams: first always 429, second always 200.
	badUp := &mockUpstream{status: 429, body: "rate limited"}
	badSrv := httptest.NewServer(http.HandlerFunc(badUp.handler))
	goodUp := &mockUpstream{status: 200, body: `{"ok":true}`}
	goodSrv := httptest.NewServer(http.HandlerFunc(goodUp.handler))
	defer badSrv.Close()
	defer goodSrv.Close()

	// Two accounts, each pointing at a different upstream via per-account BaseURL.
	accts := []config.Account{
		{ID: "idA", Name: "a", APIKey: "k1", BaseURL: badSrv.URL},
		{ID: "idB", Name: "b", APIKey: "k2", BaseURL: goodSrv.URL},
	}
	h := newTestHandler(t, "https://unused.example.com", accts, 2)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/api/version", nil)
	h.ServeHTTP(rec, req)

	if rec.Code != 200 {
		t.Fatalf("status = %d, want 200 (failover to good account)", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), `"ok":true`) {
		t.Errorf("body = %q, want good upstream response", rec.Body.String())
	}
}

func TestHandler_AllAccounts429_Returns429(t *testing.T) {
	up := &mockUpstream{status: 429, body: "rate limited"}
	srv := httptest.NewServer(http.HandlerFunc(up.handler))
	defer srv.Close()

	accts := []config.Account{
		{ID: "idA", Name: "a", APIKey: "k1", BaseURL: srv.URL},
		{ID: "idB", Name: "b", APIKey: "k2", BaseURL: srv.URL},
	}
	h := newTestHandler(t, srv.URL, accts, 2)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/api/version", nil)
	h.ServeHTTP(rec, req)

	if rec.Code != http.StatusTooManyRequests {
		t.Fatalf("status = %d, want 429", rec.Code)
	}
}

func TestHandler_SSEStreaming(t *testing.T) {
	up := &mockUpstream{
		status: 200,
		chunks: []string{"data: {\"a\":1}\n\n", "data: {\"a\":2}\n\n", "data: [DONE]\n\n"},
	}
	srv := httptest.NewServer(http.HandlerFunc(up.handler))
	defer srv.Close()

	h := newTestHandler(t, srv.URL, mkAccts(1), 1)
	rec := httptest.NewRecorder()
	req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"stream":true}`))
	h.ServeHTTP(rec, req)

	if rec.Code != 200 {
		t.Fatalf("status = %d, want 200", rec.Code)
	}
	body := rec.Body.String()
	if !strings.Contains(body, "[DONE]") {
		t.Errorf("body = %q, want to contain [DONE]", body)
	}
	if !strings.Contains(body, "\"a\":1") || !strings.Contains(body, "\"a\":2") {
		t.Errorf("body = %q, missing chunks", body)
	}
	if ct := rec.Header().Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
		t.Errorf("Content-Type = %q, want text/event-stream", ct)
	}
}

func TestHandler_RejectsUnknownPath(t *testing.T) {
	up := &mockUpstream{status: 200, body: "x"}
	srv := httptest.NewServer(http.HandlerFunc(up.handler))
	defer srv.Close()

	h := newTestHandler(t, srv.URL, mkAccts(1), 1)
	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/admin", nil)
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusNotFound {
		t.Errorf("status = %d, want 404 for unknown path", rec.Code)
	}
}

func TestHandler_NoAccounts(t *testing.T) {
	h := newTestHandler(t, "http://unused.example.com", nil, 1)
	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/api/version", nil)
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusServiceUnavailable {
		t.Errorf("status = %d, want 503", rec.Code)
	}
}

func TestHandler_StripsClientAuthorization(t *testing.T) {
	// Upstream echoes the received Authorization header in its body so we can
	// assert the proxy overrode it with the account's key.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = io.WriteString(w, `{"got_auth":"`+r.Header.Get("Authorization")+`"}`)
	}))
	defer srv.Close()

	accts := []config.Account{{ID: "idA", Name: "a", APIKey: "sk-account-key-1234", BaseURL: srv.URL}}
	h := newTestHandler(t, srv.URL, accts, 1)
	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/api/version", nil)
	req.Header.Set("Authorization", "Bearer client-should-be-stripped")
	h.ServeHTTP(rec, req)

	if !strings.Contains(rec.Body.String(), "sk-account-key-1234") {
		t.Errorf("body = %q, want account key forwarded", rec.Body.String())
	}
	if strings.Contains(rec.Body.String(), "client-should-be-stripped") {
		t.Errorf("body = %q, client Authorization leaked to upstream", rec.Body.String())
	}
}

func TestHandler_RetriesOn5xx(t *testing.T) {
	// First account returns 500, second returns 200.
	badUp := &mockUpstream{status: 500}
	badSrv := httptest.NewServer(http.HandlerFunc(badUp.handler))
	goodUp := &mockUpstream{status: 200, body: `{"ok":true}`}
	goodSrv := httptest.NewServer(http.HandlerFunc(goodUp.handler))
	defer badSrv.Close()
	defer goodSrv.Close()

	accts := []config.Account{
		{ID: "idA", Name: "a", APIKey: "k1", BaseURL: badSrv.URL},
		{ID: "idB", Name: "b", APIKey: "k2", BaseURL: goodSrv.URL},
	}
	h := newTestHandler(t, "https://unused.example.com", accts, 2)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/api/version", nil)
	h.ServeHTTP(rec, req)

	if rec.Code != 200 {
		t.Fatalf("status = %d, want 200 (failover on 500)", rec.Code)
	}
}

func TestIsAllowed(t *testing.T) {
	good := []string{
		"/api/chat", "/api/generate", "/api/tags", "/api/show", "/api/ps",
		"/api/version", "/api/delete",
		"/v1/chat/completions", "/v1/completions", "/v1/models",
		"/v1/embeddings", "/v1/files", "/v1/files/abc-123",
	}
	for _, p := range good {
		if !isAllowed(p) {
			t.Errorf("isAllowed(%q) = false, want true", p)
		}
	}
	bad := []string{"/admin", "/", "/api/unknown", "/v2/foo", "/etc/passwd"}
	for _, p := range bad {
		if isAllowed(p) {
			t.Errorf("isAllowed(%q) = true, want false", p)
		}
	}
}

// TestHandler_GoPathRoutesToGoProvider verifies /go/v1/* forces the
// opencode-go provider and strips the /go prefix before forwarding upstream.
func TestHandler_GoPathRoutesToGoProvider(t *testing.T) {
	var seenPath, seenAuth string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seenPath = r.URL.Path
		seenAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = io.WriteString(w, `{"ok":true}`)
	}))
	defer srv.Close()

	accts := []config.Account{
		{ID: "idO", Name: "o", APIKey: "kO", Provider: config.ProviderOllamaCloud, BaseURL: "http://unused-ollama"},
		{ID: "idG", Name: "g", APIKey: "kG", Provider: config.ProviderOpenCodeGo, BaseURL: srv.URL},
	}
	h := newTestHandler(t, "http://unused", accts, 2)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/go/v1/models", nil)
	h.ServeHTTP(rec, req)

	if rec.Code != 200 {
		t.Fatalf("status = %d, want 200", rec.Code)
	}
	if seenPath != "/v1/models" {
		t.Errorf("upstream path = %q, want /v1/models", seenPath)
	}
	if seenAuth != "Bearer kG" {
		t.Errorf("upstream auth = %q, want go account key", seenAuth)
	}
}

// TestHandler_ModelBasedRouting verifies a common model can be served by
// either provider and a unique Go model routes only to Go accounts.
func TestHandler_ModelBasedRouting(t *testing.T) {
	ollamaSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = io.WriteString(w, `{"from":"ollama"}`)
	}))
	goSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = io.WriteString(w, `{"from":"go"}`)
	}))
	defer ollamaSrv.Close()
	defer goSrv.Close()

	accts := []config.Account{
		{ID: "idO", Name: "o", APIKey: "kO", Provider: config.ProviderOllamaCloud, BaseURL: ollamaSrv.URL},
		{ID: "idG", Name: "g", APIKey: "kG", Provider: config.ProviderOpenCodeGo, BaseURL: goSrv.URL},
	}
	h := newTestHandler(t, "http://unused", accts, 2)

	// Unique Go model must route to the Go upstream.
	rec := httptest.NewRecorder()
	req := httptest.NewRequest("POST", "/v1/chat/completions",
		strings.NewReader(`{"model":"mimo-v2.5","messages":[]}`))
	h.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatalf("mimo status = %d", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), `"from":"go"`) {
		t.Errorf("mimo routed to wrong upstream: %s", rec.Body.String())
	}

	// Unique Ollama model must route to the Ollama upstream.
	rec = httptest.NewRecorder()
	req = httptest.NewRequest("POST", "/v1/chat/completions",
		strings.NewReader(`{"model":"gpt-oss:20b","messages":[]}`))
	h.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatalf("gpt-oss status = %d", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), `"from":"ollama"`) {
		t.Errorf("gpt-oss routed to wrong upstream: %s", rec.Body.String())
	}
}

// TestHandler_MergedModels verifies /v1/models returns the union of both
// catalogs with raw model IDs (no labels in the id field).
func TestHandler_MergedModels(t *testing.T) {
	accts := []config.Account{
		{ID: "idO", Name: "o", APIKey: "kO", Provider: config.ProviderOllamaCloud, BaseURL: "http://unused"},
		{ID: "idG", Name: "g", APIKey: "kG", Provider: config.ProviderOpenCodeGo, BaseURL: "http://unused"},
	}
	h := newTestHandler(t, "http://unused", accts, 2)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/v1/models", nil)
	h.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatalf("status = %d", rec.Code)
	}
	body := rec.Body.String()
	// Common model appears with raw ID (no label).
	if !strings.Contains(body, `"id":"glm-5"`) {
		t.Errorf("common model glm-5 missing from merged list")
	}
	// Unique Ollama model appears with raw ID (no label in id).
	if !strings.Contains(body, `"id":"gpt-oss:20b"`) {
		t.Errorf("unique ollama model missing: %s", body)
	}
	// Unique Go model appears with raw ID.
	if !strings.Contains(body, `"id":"mimo-v2.5"`) {
		t.Errorf("unique go model missing: %s", body)
	}
	// No labels should appear in any id field.
	if strings.Contains(body, `(ollama)`) || strings.Contains(body, `(go)`) {
		t.Errorf("labels should not appear in id field: %s", body)
	}
}

// TestHandler_402UsageLimitTriggersFailover verifies 402 is treated as
// retriable and triggers failover to the next account.
func TestHandler_402UsageLimitTriggersFailover(t *testing.T) {
	badSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(402)
		_, _ = io.WriteString(w, `{"error":{"message":"usage limit exhausted"}}`)
	}))
	goodSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = io.WriteString(w, `{"ok":true}`)
	}))
	defer badSrv.Close()
	defer goodSrv.Close()

	accts := []config.Account{
		{ID: "idA", Name: "a", APIKey: "k1", BaseURL: badSrv.URL},
		{ID: "idB", Name: "b", APIKey: "k2", BaseURL: goodSrv.URL},
	}
	h := newTestHandler(t, "http://unused", accts, 2)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/api/version", nil)
	h.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatalf("status = %d, want 200 after 402 failover", rec.Code)
	}
}

// TestIsRetriable and TestShouldCooldown cover the retry classification.
func TestIsRetriable(t *testing.T) {
	retriable := []int{429, 402, 500, 502, 503, 408}
	for _, s := range retriable {
		if !isRetriable(s) {
			t.Errorf("isRetriable(%d) = false, want true", s)
		}
	}
	nonRetriable := []int{200, 400, 401, 403, 404, 422}
	for _, s := range nonRetriable {
		if isRetriable(s) {
			t.Errorf("isRetriable(%d) = true, want false", s)
		}
	}
}

func TestShouldCooldown(t *testing.T) {
	if !shouldCooldown(429, "rate limit") {
		t.Error("429 should cooldown")
	}
	if !shouldCooldown(402, "usage limit") {
		t.Error("402 should cooldown")
	}
	if shouldCooldown(500, "internal error") {
		t.Error("500 should NOT cooldown")
	}
}

func TestIsBillingDisable(t *testing.T) {
	if !isBillingDisable("Insufficient credits for this account") {
		t.Error("insufficient credits should be billing disable")
	}
	if !isBillingDisable("credit balance too low") {
		t.Error("credit balance should be billing disable")
	}
	if isBillingDisable("usage limit exhausted") {
		t.Error("usage limit exhausted should NOT be billing disable")
	}
}

func TestRoutePath(t *testing.T) {
	up, prov, isGo := routePath("/go/v1/models")
	if up != "/v1/models" || prov != config.ProviderOpenCodeGo || !isGo {
		t.Errorf("routePath(/go/v1/models) = %q,%v,%v", up, prov, isGo)
	}
	up, prov, isGo = routePath("/v1/models")
	if up != "/v1/models" || prov != "" || isGo {
		t.Errorf("routePath(/v1/models) = %q,%v,%v", up, prov, isGo)
	}
}