mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 00:48:00 +00:00
feat: ModelRouter — unified OpenAI-compatible multi-source LLM gateway
- Lua adapters per upstream (transform_request/response/stream_chunk, build_headers signing hooks) - AUTO priority routing with per-model kind (chat/image), explicit source/model routing - Per-source concurrency caps with queueing, exponential backoff, AUTO failover - OpenAI-compatible API: chat completions, SSE streaming, image generations, models - Gateway key auth, web UI for adapter/source management, runtime persistence - e2e test running the real binary against mocked upstreams
This commit is contained in:
root
288
e2e/e2e_test.go
Normal file
288
e2e/e2e_test.go
Normal file
@ -0,0 +1,288 @@
|
||||
// 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"
|
||||
"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")
|
||||
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.
|
||||
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", "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 priority: 80\n")
|
||||
} else {
|
||||
sb.WriteString(" - name: " + name + "\n base_url: " + u.baseURL + "\n adapter: openai\n models:\n - id: " + name + "-m\n priority: 100\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),
|
||||
"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())
|
||||
}
|
||||
}
|
||||
|
||||
func statusOf(resp *http.Response) int {
|
||||
if resp == nil {
|
||||
return -1
|
||||
}
|
||||
return resp.StatusCode
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user