ModelRouter/e2e/e2e_test.go

// Package e2e runs the real llmsproxy binary against mocked upstreams over
// real HTTP: startup, auth, chat, streaming SSE, image generation, and
// AUTO failover when a higher-priority source fails.
package e2e

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"testing"
	"time"
)

// mockUpstream is an OpenAI-compatible upstream that can be programmed to fail.
type mockUpstream struct {
	mu      sync.Mutex
	hits    int
	fail    bool
	server  *http.Server
	baseURL string
	done    chan struct{}
}

func newMockUpstream(t *testing.T) *mockUpstream {
	m := &mockUpstream{done: make(chan struct{})}
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		m.mu.Lock()
		m.hits++
		fail := m.fail
		m.mu.Unlock()
		if fail {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(w, `{"error":"mock upstream down"}`)
			return
		}
		body, _ := io.ReadAll(r.Body)
		var req map[string]interface{}
		_ = json.Unmarshal(body, &req)
		if stream, _ := req["stream"].(bool); stream {
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(200)
			fmt.Fprintln(w, `data: {"choices":[{"delta":{"content":"Hel"}}]}`)
			fmt.Fprintln(w, `data: {"choices":[{"delta":{"content":"lo"}}]}`)
			fmt.Fprintln(w, `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`)
			fmt.Fprintln(w, "data: [DONE]")
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		fmt.Fprintf(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`)
	})
	mux.HandleFunc("/v1/images/generations", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"created":1,"data":[{"b64_json":"QUJD"}]}`)
	})
	m.server = &http.Server{Handler: mux}
	go m.server.Serve(l)
	m.baseURL = "http://" + l.Addr().String()
	t.Cleanup(func() {
		close(m.done)
		_ = m.server.Close()
	})
	return m
}

func (m *mockUpstream) Hits() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.hits
}

func (m *mockUpstream) SetFail(f bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.fail = f
}

// gatewayUnderTest is the real binary, started from a temp config.
type gatewayUnderTest struct {
	cmd  *exec.Cmd
	addr string
	key  string
	log  bytes.Buffer
}

func buildBinary(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	bin := filepath.Join(dir, "llmsproxy")
	if runtime.GOOS == "windows" {
		bin += ".exe" // go build writes the exact name; Windows needs the suffix to exec
	}
	out, err := exec.Command("go", "build", "-tags", "luajit", "-o", bin, "llmsproxy/cmd/llmsproxy").CombinedOutput()
	if err != nil && runtime.GOOS == "windows" {
		// local Windows dev box may lack LuaJIT: fall back to the bundled
		// Lua runtime (the e2e adapters used here are passthrough-only)
		out, err = exec.Command("go", "build", "-o", bin, "llmsproxy/cmd/llmsproxy").CombinedOutput()
	}
	if err != nil {
		t.Fatalf("build: %v\n%s", err, out)
	}
	return bin
}

func startGateway(t *testing.T, bin, listen, cfgPath string) *gatewayUnderTest {
	t.Helper()
	g := &gatewayUnderTest{addr: listen, key: "sk-e2e-0001"}
	g.cmd = exec.Command(bin, "-config", cfgPath)
	g.cmd.Stdout = &g.log
	g.cmd.Stderr = &g.log
	if err := g.cmd.Start(); err != nil {
		t.Fatalf("start: %v", err)
	}
	t.Cleanup(func() {
		if g.cmd.Process != nil {
			_ = g.cmd.Process.Kill()
			_, _ = g.cmd.Process.Wait()
		}
	})
	// wait for readiness
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		req, _ := http.NewRequest("GET", "http://"+g.addr+"/v1/models", nil)
		req.Header.Set("Authorization", "Bearer "+g.key)
		resp, err := http.DefaultClient.Do(req)
		if err == nil {
			io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				return g
			}
		}
		if g.cmd.ProcessState != nil {
			t.Fatalf("gateway exited early:\n%s", g.log.String())
		}
		time.Sleep(100 * time.Millisecond)
	}
	t.Fatalf("gateway did not become ready:\n%s", g.log.String())
	return nil
}

func (g *gatewayUnderTest) do(method, path string, body string, authed bool) (*http.Response, string) {
	var r io.Reader
	if body != "" {
		r = strings.NewReader(body)
	}
	req, _ := http.NewRequest(method, "http://"+g.addr+path, r)
	if authed {
		req.Header.Set("Authorization", "Bearer "+g.key)
	}
	if body != "" {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, ""
	}
	data, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	return resp, string(data)
}

// writeConfig writes a temp gateway config pointing at the mock upstreams,
// listening on the given address. Every non-image upstream becomes a chat
// source; the first named one ("good") is the highest-priority AUTO slot.
func writeConfig(t *testing.T, dir, listenAddr string, upstreams map[string]*mockUpstream) string {
	t.Helper()
	var sb strings.Builder
	sb.WriteString("listen: " + listenAddr + "\n")
	sb.WriteString("gateway_keys:\n  - sk-e2e-0001\n")
	sb.WriteString("default_model: AUTO\n")
	sb.WriteString("adapter_dir: " + filepath.Join(dir, "adapters") + "\n")
	sb.WriteString("runtime_file: " + filepath.Join(dir, "runtime.json") + "\n")
	sb.WriteString("sources:\n")
	order := []string{"good", "fallback", "image"}
	prio := map[string]int{"good": 100, "fallback": 50}
	for _, name := range order {
		u, ok := upstreams[name]
		if !ok {
			continue
		}
		if name == "image" {
			sb.WriteString("  - name: imagegen\n    base_url: " + u.baseURL + "\n    adapter: openai\n    models:\n      - id: flux-1\n        kind: image\n        priority: 80\n")
		} else {
			sb.WriteString("  - name: " + name + "\n    base_url: " + u.baseURL + "\n    adapter: openai\n    models:\n      - id: " + name + "-m\n        priority: " + fmt.Sprint(prio[name]) + "\n")
		}
	}
	path := filepath.Join(dir, "config.yaml")
	if err := os.WriteFile(path, []byte(sb.String()), 0o644); err != nil {
		t.Fatalf("write config: %v", err)
	}
	return path
}

func TestEndToEnd(t *testing.T) {
	upstreams := map[string]*mockUpstream{
		"good":     newMockUpstream(t),
		"fallback": newMockUpstream(t),
		"image":    newMockUpstream(t),
	}

	dir := t.TempDir()

	// pick a fixed free port for the gateway
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("pick port: %v", err)
	}
	addr := l.Addr().String()
	l.Close()

	cfg := writeConfig(t, dir, addr, upstreams)

	bin := buildBinary(t)
	g := startGateway(t, bin, addr, cfg)

	// 1. no key -> 401
	resp, body := g.do("GET", "/v1/models", "", false)
	if resp == nil || resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("unauth status = %v body=%q", statusOf(resp), body)
	}

	// 2. list models
	resp, body = g.do("GET", "/v1/models", "", true)
	if resp == nil || resp.StatusCode != 200 {
		t.Fatalf("models status=%v body=%q", statusOf(resp), body)
	}
	if !strings.Contains(body, "good-m") || !strings.Contains(body, "flux-1") {
		t.Fatalf("models missing entries: %s", body)
	}

	// 3. single chat (explicit model)
	resp, body = g.do("POST", "/v1/chat/completions",
		`{"model":"good-m","messages":[{"role":"user","content":"hi"}]}`, true)
	if resp == nil || resp.StatusCode != 200 {
		t.Fatalf("chat status=%v body=%q", statusOf(resp), body)
	}
	var chat map[string]interface{}
	_ = json.Unmarshal([]byte(body), &chat)
	msg := chat["choices"].([]interface{})[0].(map[string]interface{})["message"].(map[string]interface{})
	if msg["content"] != "pong" {
		t.Fatalf("chat content = %q", msg["content"])
	}

	// 4. streaming SSE
	resp, body = g.do("POST", "/v1/chat/completions",
		`{"model":"good-m","stream":true,"messages":[{"role":"user","content":"hi"}]}`, true)
	if resp == nil || resp.StatusCode != 200 {
		t.Fatalf("stream status=%v body=%q", statusOf(resp), body)
	}
	if !strings.Contains(body, "Hel") || !strings.Contains(body, "lo") || !strings.Contains(body, "[DONE]") {
		t.Fatalf("stream body = %q", body)
	}

	// 5. image generation
	resp, body = g.do("POST", "/v1/images/generations",
		`{"model":"flux-1","prompt":"a cat"}`, true)
	if resp == nil || resp.StatusCode != 200 {
		t.Fatalf("image status=%v body=%q", statusOf(resp), body)
	}
	if !strings.Contains(body, "QUJD") {
		t.Fatalf("image body = %q", body)
	}

	// 6. AUTO failover: make the high-priority source fail, chat should still work
	before := upstreams["good"].Hits()
	upstreams["good"].SetFail(true)
	resp, body = g.do("POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`, true)
	upstreams["good"].SetFail(false)
	if resp == nil || resp.StatusCode != 200 {
		t.Fatalf("failover status=%v body=%q", statusOf(resp), body)
	}
	// the failing upstream must have received the attempt (proving fallback happened)
	if upstreams["good"].Hits() <= before {
		t.Fatalf("failover did not try the failing provider (hits %d -> %d)", before, upstreams["good"].Hits())
	}
}

// TestEndToEndAuto503: with every AUTO slot failing, the gateway must answer
// 503 whose message summarizes each failed tier/source/model instead of a
// bare "no provider available" (P8).
func TestEndToEndAuto503(t *testing.T) {
	upstreams := map[string]*mockUpstream{
		"good": newMockUpstream(t),
	}
	dir := t.TempDir()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("pick port: %v", err)
	}
	addr := l.Addr().String()
	l.Close()
	cfg := writeConfig(t, dir, addr, upstreams)
	bin := buildBinary(t)
	g := startGateway(t, bin, addr, cfg)

	upstreams["good"].SetFail(true)
	resp, body := g.do("POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`, true)
	if resp == nil || resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("status=%v body=%q", statusOf(resp), body)
	}
	if !strings.Contains(body, "auto providers failed") || !strings.Contains(body, "good/good-m") {
		t.Fatalf("503 must summarize the failed slot, body=%q", body)
	}
}

func statusOf(resp *http.Response) int {
	if resp == nil {
		return -1
	}
	return resp.StatusCode
}