ollama-proxy/internal/cli/integration_test.go

package cli

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// findBinary locates the ollama-proxy executable built by `go test`. The test
// binary is in the package dir; the project binary is one level up. We try both
// and fall back to building one on the fly.
func findBinary(t *testing.T) string {
	t.Helper()
	candidates := []string{
		filepath.Join("..", "..", "ollama-proxy"),
		"ollama-proxy",
	}
	for _, c := range candidates {
		if abs, err := filepath.Abs(c); err == nil {
			if fi, err := os.Stat(abs); err == nil && !fi.IsDir() {
				return abs
			}
		}
	}
	// Build it now into a temp path.
	out := filepath.Join(t.TempDir(), "ollama-proxy")
	cmd := exec.Command("go", "build", "-o", out, ".")
	cmd.Dir = filepath.Join("..", "..")
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("go build: %v\n%s", err, out)
	}
	return out
}

// freePort picks an unused TCP port for the proxy to listen on.
func freePort(t *testing.T) string {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()
	return l.Addr().String()
}

// startProxy starts the binary as a subprocess and returns a cancel function
// that stops it.
func startProxy(t *testing.T, bin, xdgDir, addr, upstream string) func() {
	t.Helper()
	cmd := exec.Command(bin, "serve",
		"--addr", addr,
		"--base-url", upstream,
		"--cooldown", "1s",
		"--retries", "3",
		"--log-level", "debug",
	)
	cmd.Env = append(os.Environ(),
		"XDG_CONFIG_HOME="+xdgDir,
		"OLLAMA_PROXY_LOG_LEVEL=debug",
	)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Start(); err != nil {
		t.Fatalf("start proxy: %v", err)
	}
	// Wait for the port to accept connections.
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		c, err := net.Dial("tcp", addr)
		if err == nil {
			c.Close()
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	return func() {
		_ = cmd.Process.Signal(os.Interrupt)
		done := make(chan struct{})
		go func() { _ = cmd.Wait(); close(done) }()
		select {
		case <-done:
		case <-time.After(2 * time.Second):
			_ = cmd.Process.Kill()
			<-done
		}
		t.Logf("proxy stdout: %s", stdout.String())
		t.Logf("proxy stderr: %s", stderr.String())
	}
}

// TestEndToEnd_ProxyForwardsToUpstream builds the binary and verifies the full
// stack: account JSON → balancer → handler → real HTTP server → upstream.
func TestEndToEnd_ProxyForwardsToUpstream(t *testing.T) {
	// Upstream that echoes the auth and path.
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = fmt.Fprintf(w, `{"path":%q,"auth":%q}`, r.URL.Path, r.Header.Get("Authorization"))
	}))
	defer upstream.Close()

	bin := findBinary(t)
	xdg := t.TempDir()
	// Write accounts.json directly via the binary's `accounts add`.
	for _, a := range []struct{ key, name string }{
		{"sk-keyAAAAAAAA", "a1"},
		{"sk-keyBBBBBBBB", "a2"},
	} {
		cmd := exec.Command(bin, "accounts", "add", a.key, "--name", a.name)
		cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
		if out, err := cmd.CombinedOutput(); err != nil {
			t.Fatalf("accounts add %s: %v\n%s", a.name, err, out)
		}
	}

	addr := freePort(t)
	stop := startProxy(t, bin, xdg, addr, upstream.URL)
	defer stop()

	resp, err := http.Get("http://" + addr + "/api/version")
	if err != nil {
		t.Fatalf("GET /api/version: %v", err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		t.Fatalf("status = %d, body = %s", resp.StatusCode, body)
	}
	if !strings.Contains(string(body), "sk-key") {
		t.Errorf("body = %s, want account key forwarded", body)
	}
	if !strings.Contains(string(body), "/api/version") {
		t.Errorf("body = %s, want path echoed", body)
	}
}

// TestEndToEnd_SSEStreaming verifies SSE passthrough through the real binary.
func TestEndToEnd_SSEStreaming(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Transfer-Encoding", "chunked")
		w.WriteHeader(200)
		flusher, _ := w.(http.Flusher)
		for i := 0; i < 3; i++ {
			_, _ = io.WriteString(w, "data: {\"i\":"+string(rune('0'+i))+"}\n\n")
			if flusher != nil {
				flusher.Flush()
			}
		}
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
		if flusher != nil {
			flusher.Flush()
		}
	}))
	defer upstream.Close()

	bin := findBinary(t)
	xdg := t.TempDir()
	cmd := exec.Command(bin, "accounts", "add", "sk-testkey", "--name", "a1")
	cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("accounts add: %v\n%s", err, out)
	}

	addr := freePort(t)
	stop := startProxy(t, bin, xdg, addr, upstream.URL)
	defer stop()

	resp, err := http.Post("http://"+addr+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`))
	if err != nil {
		t.Fatalf("POST: %v", err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if !strings.Contains(string(body), "[DONE]") {
		t.Errorf("body = %s, want [DONE]", body)
	}
	if !strings.Contains(string(body), `"i":0`) || !strings.Contains(string(body), `"i":1`) {
		t.Errorf("body = %s, missing chunks", body)
	}
}

// TestEndToEnd_FailoverOn429 verifies the binary fails over between accounts.
func TestEndToEnd_FailoverOn429(t *testing.T) {
	// Upstream that 429s for the first account key and 200s for the second.
	var firstSeen string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization")
		if firstSeen == "" {
			firstSeen = auth
		}
		if auth == "Bearer sk-bad" {
			http.Error(w, "rate limited", 429)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = io.WriteString(w, `{"ok":true}`)
	}))
	defer upstream.Close()

	bin := findBinary(t)
	xdg := t.TempDir()
	for _, a := range []struct{ key, name string }{
		{"sk-bad", "bad"},
		{"sk-good", "good"},
	} {
		cmd := exec.Command(bin, "accounts", "add", a.key, "--name", a.name)
		cmd.Env = append(os.Environ(), "XDG_CONFIG_HOME="+xdg)
		if out, err := cmd.CombinedOutput(); err != nil {
			t.Fatalf("accounts add %s: %v\n%s", a.name, err, out)
		}
	}

	addr := freePort(t)
	stop := startProxy(t, bin, xdg, addr, upstream.URL)
	defer stop()

	resp, err := http.Get("http://" + addr + "/api/version")
	if err != nil {
		t.Fatalf("GET: %v", err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		t.Fatalf("status = %d, body = %s (want 200 after failover)", resp.StatusCode, body)
	}
	if !strings.Contains(string(body), `"ok":true`) {
		t.Errorf("body = %s, want good response", body)
	}
}