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()
	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)
		}
	}
}