ModelRouter/e2e/e2e_test.go

// Package e2e runs the real modelrouter 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, "modelrouter")
	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, "modelrouter/cmd/modelrouter").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, "modelrouter/cmd/modelrouter").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"}
	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")
		} else {
			sb.WriteString("  - name: " + name + "\n    base_url: " + u.baseURL + "\n    adapter: openai\n    models:\n      - id: " + name + "-m\n")
		}
	}
	// AUTO 链现在是显式配置（25970a2 移除了 models.priority）。tier 越小越先试。
	sb.WriteString("auto:\n  - model: good-m\n    source: good\n    tier: 1\n")
	sb.WriteString("  - model: fallback-m\n    source: fallback\n    tier: 2\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
}

// The startup warning about key configuration must describe the key set that
// AUTH actually uses (cfg.Keys, via core.ListKeys), not the legacy
// gateway_keys YAML list.
//
// It used to check GatewayKeys() only. seedKeys copies gateway_keys into keys[]
// on first start and the YAML list stops being consulted, so once that list was
// emptied — e.g. after rotating the starter key away — the process logged
// "every request will be rejected" while seven working keys, one of them admin,
// were in service. Observed in production: the warning appeared and that same
// key returned HTTP 200 on /v1/models.
//
// The fixture reproduces exactly that shape: gateway_keys EMPTY, with the
// working key present only in keys[].
func TestStartupWarningReflectsRealKeysNotLegacyList(t *testing.T) {
	up := newMockUpstream(t)
	dir := t.TempDir()

	// Seed a real install once so keys[] gets populated from gateway_keys, then
	// rewrite the file to look like a rotated install: legacy list cleared, the
	// seeded key still present in keys[].
	path := writeConfig(t, dir, "127.0.0.1:0", map[string]*mockUpstream{"good": up})
	bin := buildBinary(t)

	// First run creates the sealed config with keys[] populated.
	seedPort := freePort(t)
	cfg1 := writeConfig(t, dir, "127.0.0.1:"+seedPort, map[string]*mockUpstream{"good": up})
	g1 := startGateway(t, bin, "127.0.0.1:"+seedPort, cfg1)
	g1.stop()

	data, err := os.ReadFile(cfg1)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(data), "keys:") {
		t.Fatalf("first run did not populate keys[]:\n%s", data)
	}

	// Now clear the legacy gateway_keys list, keeping keys[] intact. The
	// rewritten file may use different indentation, so match loosely.
	cleared, n := clearGatewayKeys(string(data))
	if n == 0 {
		t.Fatalf("could not find a gateway_keys entry to clear in:\n%s", data)
	}
	if err := os.WriteFile(cfg1, []byte(cleared), 0o644); err != nil {
		t.Fatal(err)
	}

	// Start again and inspect the log. The config pins its listen address, so
	// rewrite that field to the port we will actually poll.
	port2 := freePort(t)
	rehosted, ok := replaceListen(string(cleared), "127.0.0.1:"+port2)
	if !ok {
		t.Fatalf("could not rewrite listen in:\n%s", cleared)
	}
	cfg2 := filepath.Join(dir, "config2.yaml")
	if err := os.WriteFile(cfg2, []byte(rehosted), 0o644); err != nil {
		t.Fatal(err)
	}
	g2 := startGateway(t, bin, "127.0.0.1:"+port2, cfg2)
	defer g2.stop()

	// Sanity: the key really does authenticate, so the warning would be a lie.
	resp, _ := g2.do("GET", "/v1/models", "", true)
	if resp == nil || resp.StatusCode != http.StatusOK {
		t.Fatalf("key from keys[] did not authenticate (status %v); test fixture is wrong",
			resp)
	}

	logs := g2.logString()
	if strings.Contains(logs, "gateway_keys is EMPTY") {
		t.Errorf("startup logged 'gateway_keys is EMPTY ... every request is rejected' while "+
			"a working admin key was in service — the warning must read the authoritative "+
			"key set (core.ListKeys), not the legacy YAML list.\nlog:\n%s", logs)
	}
	_ = path
	_ = g1
}

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()
	_, port, _ := net.SplitHostPort(l.Addr().String())
	return port
}

// stop terminates the gateway and waits, so its log buffer is complete.
func (g *gatewayUnderTest) stop() {
	if g == nil || g.cmd == nil || g.cmd.Process == nil {
		return
	}
	_ = g.cmd.Process.Kill()
	_, _ = g.cmd.Process.Wait()
}

// logString returns everything the gateway has written so far.
func (g *gatewayUnderTest) logString() string { return g.log.String() }

// clearGatewayKeys empties the gateway_keys list in a config, tolerating any
// indentation the YAML writer chose. Returns the new text and how many entries
// it removed.
func clearGatewayKeys(src string) (string, int) {
	lines := strings.Split(src, "\n")
	var out []string
	removed := 0
	inList := false
	for _, l := range lines {
		trimmed := strings.TrimSpace(l)
		if strings.HasPrefix(l, "gateway_keys:") {
			out = append(out, "gateway_keys: []")
			inList = true
			removed++
			continue
		}
		if inList {
			if strings.HasPrefix(trimmed, "- ") {
				removed++
				continue
			}
			inList = false
		}
		out = append(out, l)
	}
	return strings.Join(out, "\n"), removed
}

// replaceListen rewrites the top-level listen: value.
func replaceListen(src, addr string) (string, bool) {
	lines := strings.Split(src, "\n")
	for i, l := range lines {
		if strings.HasPrefix(l, "listen:") {
			lines[i] = "listen: " + addr
			return strings.Join(lines, "\n"), true
		}
	}
	return src, false
}