package gateway import ( "encoding/json" "fmt" "io" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" "llmsproxy/internal/config" "llmsproxy/internal/core" ) func mockUpstream() *httptest.Server { return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 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}}`) })) } func newTestGateway(t *testing.T, srcs ...config.Source) *Gateway { t.Helper() td := t.TempDir() cfgPath := filepath.Join(td, "config.yaml") os.WriteFile(cfgPath, []byte("listen: :0"), 0644) cfg := &config.Config{ Path: cfgPath, AdapterDir: filepath.Join(td, "adapters"), RuntimeFile: filepath.Join(td, "runtime.json"), GatewayKeys: []string{"sk-test"}, Sources: srcs, } if err := cfg.ApplyDefaults(); err != nil { t.Fatal(err) } c, err := core.NewFromConfig(cfg) if err != nil { t.Fatalf("core: %v", err) } t.Cleanup(c.Close) g, err := New(c, []string{"sk-test"}) if err != nil { t.Fatalf("gateway: %v", err) } return g } func doReq(t *testing.T, g *Gateway, method, path, body string) *httptest.ResponseRecorder { t.Helper() req, _ := http.NewRequest(method, path, strings.NewReader(body)) req.Header.Set("Authorization", "Bearer sk-test") if body != "" { req.Header.Set("Content-Type", "application/json") } rr := httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) return rr } // upstreamCtrl toggles a mocked upstream's behavior between requests. type upstreamCtrl struct { status int // 0 = healthy; else every request fails with that status hits int // chat call count } // upstream returns a mocked OpenAI upstream driven by ctrl.status. func upstream(t *testing.T, ctrl *upstreamCtrl) *httptest.Server { t.Helper() return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { ctrl.hits++ if ctrl.status != 0 { w.WriteHeader(ctrl.status) fmt.Fprint(w, `{"error":"boom"}`) return } w.Header().Set("Content-Type", "application/json") fmt.Fprintf(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`) })) } // TestChatAutoChainTierFailover: AUTO chain, first slot hard-fails, the pass // moves on within the same tier and the request is served by the next slot. func TestChatAutoChainTierFailover(t *testing.T) { a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{} aUp := upstream(t, a) bUp := upstream(t, b) defer aUp.Close() defer bUp.Close() g := newTestGateway(t, config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}}, config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m", Priority: 10}}}, ) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String()) } var cc ChatCompletion _ = json.Unmarshal(rr.Body.Bytes(), &cc) if cc.Model != "b-m" { t.Fatalf("AUTO served %q, want b-m", cc.Model) } if a.hits == 0 || b.hits == 0 { t.Fatalf("hit counts a=%d b=%d, want both > 0", a.hits, b.hits) } } // TestChatAutoChain503Summary: every AUTO slot fails -> 503 whose message // names each failed tier/source/model. func TestChatAutoChain503Summary(t *testing.T) { a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{status: 500} aUp := upstream(t, a) bUp := upstream(t, b) defer aUp.Close() defer bUp.Close() g := newTestGateway(t, config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}}, config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m", Priority: 10}}}, ) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != http.StatusServiceUnavailable { t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String()) } if !strings.Contains(rr.Body.String(), "auto providers failed") || !strings.Contains(rr.Body.String(), "a/a-m") || !strings.Contains(rr.Body.String(), "b/b-m") { t.Fatalf("503 must summarize every tier, body=%s", rr.Body.String()) } } // TestChatAutoQuotaSkip: a slot whose token quota is exhausted is dropped // from scheduling; with no other slot the chain answers 503 naming the quota. func TestChatAutoQuotaSkip(t *testing.T) { ctrl := &upstreamCtrl{} up := upstream(t, ctrl) defer up.Close() g := newTestGateway(t, config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}}, ) // one slot with an hourly quota of 1 token rr := doReq(t, g, "PUT", "/api/auto", `{"rules":[{"model":"a-m","tier":0,"token_quota":1,"period":"hour"}]}`) if rr.Code != 200 { t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String()) } // first request consumes 4 tokens -> quota exhausted rr = doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("first status=%d body=%s", rr.Code, rr.Body.String()) } // second request must skip the exhausted slot and fail 503 rr = doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != http.StatusServiceUnavailable { t.Fatalf("quota status=%d body=%s", rr.Code, rr.Body.String()) } if !strings.Contains(rr.Body.String(), "quota exhausted") { t.Fatalf("503 must name the quota reason, body=%s", rr.Body.String()) } if ctrl.hits != 1 { t.Fatalf("upstream hits = %d, want 1 (exhausted slot must not be called)", ctrl.hits) } } // TestAutoStatesReportChainHealth: GET /api/auto reports per-slot health for // the priority-page UI; a chain edit resets the failure state to zero. func TestAutoStatesReportChainHealth(t *testing.T) { ctrl := &upstreamCtrl{status: 500} up := upstream(t, ctrl) defer up.Close() g := newTestGateway(t, config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}}, ) doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) fetch := func() []core.AutoSlotState { rr := doReq(t, g, "GET", "/api/auto", "") if rr.Code != 200 { t.Fatalf("get auto status=%d body=%s", rr.Code, rr.Body.String()) } var body struct { Rules []config.ModelScope `json:"rules"` States []core.AutoSlotState `json:"states"` } if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil { t.Fatalf("unmarshal: %v", err) } return body.States } st := fetch() if len(st) != 1 || st[0].Model != "a-m" || st[0].Source != "a" { t.Fatalf("want 1 slot a/a-m, got %#v", st) } if st[0].FailCount == 0 || !st[0].Cooling { t.Fatalf("slot must report the failure (fail=%d cooling=%v)", st[0].FailCount, st[0].Cooling) } doReq(t, g, "PUT", "/api/auto", `{"rules":[{"model":"a-m","tier":0}]}`) st = fetch() if st[0].FailCount != 0 || st[0].Cooling { t.Fatalf("edit must reset health, got %#v", st[0]) } } // TestAutoSaveResetsCooldown: editing the AUTO chain clears the cooldown of // its slots, so a fixed upstream is schedulable again without waiting (P1). func TestAutoSaveResetsCooldown(t *testing.T) { ctrl := &upstreamCtrl{status: 500} up := upstream(t, ctrl) defer up.Close() g := newTestGateway(t, config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}}, ) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != http.StatusServiceUnavailable { t.Fatalf("expect 503 while upstream down, got %d", rr.Code) } p := g.core.ProviderForSlot("a-m", "a") if p == nil || p.ModelAvailable("a-m") { t.Fatal("a-m must be cooling after the failure") } // editing the chain (same rules) must clear the cooldown immediately rr = doReq(t, g, "PUT", "/api/auto", `{"rules":[{"model":"a-m","tier":0}]}`) if rr.Code != 200 { t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String()) } if !p.ModelAvailable("a-m") { t.Fatal("SaveAutoRules must reset the slot cooldown") } // healed upstream -> AUTO serves again on the next request ctrl.status = 0 rr = doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("AUTO after reset status=%d body=%s", rr.Code, rr.Body.String()) } } func TestChatSingle(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"mock-model","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String()) } var cc ChatCompletion if err := json.Unmarshal(rr.Body.Bytes(), &cc); err != nil { t.Fatalf("unmarshal: %v", err) } if cc.Choices[0].Message.Content != "pong" { t.Fatalf("content = %q", cc.Choices[0].Message.Content) } if cc.Usage == nil || cc.Usage.Total != 4 { t.Fatalf("usage = %+v", cc.Usage) } if cc.Model != "mock-model" { t.Fatalf("model = %q", cc.Model) } } func TestChatDisableThinkingPassthrough(t *testing.T) { got := "" up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { b, _ := io.ReadAll(r.Body) got = string(b) w.Header().Set("Content-Type", "application/json") fmt.Fprint(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}]}`) })) defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"mock-model","disable_thinking":true,"messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String()) } var sent map[string]interface{} if err := json.Unmarshal([]byte(got), &sent); err != nil { t.Fatalf("upstream body: %v", err) } if _, has := sent["disable_thinking"]; has { t.Fatalf("disable_thinking not stripped: %s", got) } // openai adapter strips disable_thinking; deepseek would map it to extra_body.thinking. // with the passthrough fix the flag now reaches the VM at all. } func TestChatMultimodalPassthrough(t *testing.T) { gotBody := "" up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { b, _ := io.ReadAll(r.Body) gotBody = string(b) w.Header().Set("Content-Type", "application/json") fmt.Fprint(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}]}`) })) defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"mock-model","messages":[{"role":"user","content":[ {"type":"text","text":"what is this?"}, {"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}} ]}]}`) if rr.Code != 200 { t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String()) } var sent struct { Messages []struct { Content []map[string]interface{} `json:"content"` } `json:"messages"` } if err := json.Unmarshal([]byte(gotBody), &sent); err != nil { t.Fatalf("upstream body: %v", err) } if len(sent.Messages) != 1 || len(sent.Messages[0].Content) != 2 { t.Fatalf("multimodal content lost: %s", gotBody) } } func TestChatAUTO(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "low", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "low-m", Priority: 10}}}, config.Source{Name: "high", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "high-m", Priority: 100}}}, ) // no model -> AUTO -> picks the highest priority source rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String()) } var cc ChatCompletion _ = json.Unmarshal(rr.Body.Bytes(), &cc) if cc.Model != "high-m" { t.Fatalf("AUTO picked %q, want high-m", cc.Model) } } func TestChatAuthRequired(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) req, _ := http.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"messages":[{"role":"user","content":"hi"}]}`)) rr := httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d", rr.Code) } } func TestChatStream(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"mock-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`) body := rr.Body.String() if !strings.Contains(body, "data: [DONE]") { t.Fatalf("missing DONE, body=%s", body) } if !strings.Contains(body, "Hel") || !strings.Contains(body, "lo") { t.Fatalf("missing content chunks, body=%s", body) } } func TestImageGeneration(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") fmt.Fprint(w, `{"created":1,"data":[{"b64_json":"QUJD"}]}`) })) defer up.Close() g := newTestGateway(t, config.Source{Name: "img", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "flux-1", Kind: "image"}}}) rr := doReq(t, g, "POST", "/v1/images/generations", `{"model":"flux-1","prompt":"a cat"}`) if rr.Code != 200 { t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String()) } var out map[string]interface{} _ = json.Unmarshal(rr.Body.Bytes(), &out) data, _ := out["data"].([]interface{}) if len(data) != 1 { t.Fatalf("image data len = %d", len(data)) } } func TestImageAutoFallsOnlyToImageProviders(t *testing.T) { imageHits := 0 chatHits := 0 img := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { imageHits++ w.Header().Set("Content-Type", "application/json") fmt.Fprint(w, `{"created":1,"data":[{"b64_json":"QUJD"}]}`) })) defer img.Close() chatUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { chatHits++ w.Header().Set("Content-Type", "application/json") fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`) })) defer chatUp.Close() g := newTestGateway(t, config.Source{Name: "chat", BaseURL: chatUp.URL, Adapter: "openai", Models: []config.Model{{ID: "chat-m", Priority: 100}}}, config.Source{Name: "img", BaseURL: img.URL, Adapter: "openai", Models: []config.Model{{ID: "flux", Kind: "image", Priority: 1}}}, ) rr := doReq(t, g, "POST", "/v1/images/generations", `{"model":"AUTO","prompt":"a cat"}`) if rr.Code != 200 { t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String()) } if chatHits != 0 { t.Fatalf("image AUTO hit chat-only provider: %d chat hits", chatHits) } if imageHits == 0 { t.Fatalf("image AUTO did not hit image provider") } } func TestKimicodeSigning(t *testing.T) { var gotAuth, gotSign string up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.ReadAll(r.Body) gotAuth = r.Header.Get("Authorization") gotSign = r.Header.Get("X-App-Sign") w.Header().Set("Content-Type", "application/json") fmt.Fprintf(w, `{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`) })) defer up.Close() g := newTestGateway(t, config.Source{ Name: "kimi", BaseURL: up.URL, Adapter: "kimicode", APIKey: "sk-kimi", Models: []config.Model{{ID: "kimi-k2"}}, Meta: map[string]interface{}{"app_id": "app-1", "app_secret": "s3cr3t", "app_agent": "code-agent", "api_key": "sk-kimi"}, }) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"kimi-k2","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String()) } if gotAuth == "" || !strings.Contains(gotAuth, "sk-kimi") { t.Fatalf("expected signed auth, got %q", gotAuth) } if gotSign == "" { t.Fatalf("expected app signature header") } } func TestModelRoutingPrefix(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-a"}}}, config.Source{Name: "b", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-b"}}}, ) rr := doReq(t, g, "POST", "/v1/chat/completions", `{"model":"model-b","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("status=%d", rr.Code) } } func TestModelsEndpoint(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-a"}}}, config.Source{Name: "b", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-b"}}}, ) rr := doReq(t, g, "GET", "/v1/models", "") if rr.Code != 200 { t.Fatalf("status=%d", rr.Code) } var out map[string]interface{} _ = json.Unmarshal(rr.Body.Bytes(), &out) if !strings.Contains(rr.Body.String(), "model-a") || !strings.Contains(rr.Body.String(), "model-b") { t.Fatalf("missing models: %s", rr.Body.String()) } } func TestWebUIServesPage(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) rr := doReq(t, g, "GET", "/", "") if rr.Code != 200 { t.Fatalf("status=%d", rr.Code) } if !strings.Contains(rr.Body.String(), "ModelRouter") { t.Fatalf("ui not served") } } func TestAdaptersAPIUpload(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) // upload adapter rr := doReq(t, g, "POST", "/api/adapters", `{"name":"testadp","code":"return {name='testadp',endpoint='/chat/completions',transform_request=function(raw) return raw end,transform_response=function(raw) return raw end}"}`) if rr.Code != 200 { t.Fatalf("upload status=%d body=%s", rr.Code, rr.Body.String()) } rr = doReq(t, g, "GET", "/api/status", "") if !strings.Contains(rr.Body.String(), "testadp") { t.Fatalf("adapter not listed: %s", rr.Body.String()) } } func TestSourcesAPIAddAndPersist(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t) rr := doReq(t, g, "POST", "/api/sources", fmt.Sprintf(`{"name":"added","base_url":"%s","adapter":"openai","models":[{"id":"new-m","priority":5}]}`, up.URL)) if rr.Code != 200 { t.Fatalf("add source status=%d body=%s", rr.Code, rr.Body.String()) } rr = doReq(t, g, "GET", "/v1/models", "") if !strings.Contains(rr.Body.String(), "new-m") { t.Fatalf("new model not live: %s", rr.Body.String()) } // verify source persisted to config.yaml if _, err := os.Stat(g.core.Config().Path); err != nil { t.Fatalf("config file not written: %v", err) } } func TestUIRequiresAuth(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) // / without a key -> redirect to /login req, _ := http.NewRequest("GET", "/", nil) rr := httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusFound { t.Fatalf("expected 302 for /, got %d", rr.Code) } if loc := rr.Header().Get("Location"); !strings.Contains(loc, "/login") { t.Fatalf("expected redirect to /login, got %q", loc) } // /api/status without a key -> 401 req, _ = http.NewRequest("GET", "/api/status", nil) rr = httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusUnauthorized { t.Fatalf("expected 401 for /api/status, got %d", rr.Code) } // /login page is public req, _ = http.NewRequest("GET", "/login", nil) rr = httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusOK || !strings.Contains(rr.Body.String(), "登录") { t.Fatalf("login page: %d %s", rr.Code, rr.Body.String()) } } func TestLoginAPIAndCookie(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) // bad key -> 401 req, _ := http.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"wrong"}`)) req.Header.Set("Content-Type", "application/json") rr := httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusUnauthorized { t.Fatalf("bad login: %d", rr.Code) } // good key -> cookie req, _ = http.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"sk-test"}`)) req.Header.Set("Content-Type", "application/json") rr = httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusOK { t.Fatalf("login: %d", rr.Code) } cookies := rr.Result().Cookies() if len(cookies) == 0 || cookies[0].Name != "gw_key" { t.Fatalf("no gw_key cookie set") } // use cookie to access /api/status req, _ = http.NewRequest("GET", "/api/status", nil) req.AddCookie(cookies[0]) rr = httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusOK { t.Fatalf("cookie authed status: %d", rr.Code) } var out map[string]interface{} _ = json.Unmarshal(rr.Body.Bytes(), &out) if out["base_url"] == "" { t.Fatalf("status missing base_url: %s", rr.Body.String()) } if !strings.Contains(rr.Body.String(), "sk-test") { t.Fatalf("status missing gateway_keys: %s", rr.Body.String()) } } func TestAPIChatInternal(t *testing.T) { up := mockUpstream() defer up.Close() g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}}) rr := doReq(t, g, "POST", "/api/chat", `{"model":"mock-model","messages":[{"role":"user","content":"hi"}]}`) if rr.Code != 200 { t.Fatalf("api chat status=%d body=%s", rr.Code, rr.Body.String()) } if !strings.Contains(rr.Body.String(), "pong") { t.Fatalf("api chat body=%s", rr.Body.String()) } } func TestHasScopeModelWithSourcePrefix(t *testing.T) { // zen serves deepseek-v4-flash-free, so the pinning prefix strips; deepseek // serves deepseek-v4-pro but NOT v4-flash-free, so a model id like // "deepseek-v4-flash-free" (source deepseek + model id with dashes) must NOT // be corrupted by prefix stripping. g := newTestGateway(t, config.Source{Name: "deepseek", BaseURL: "http://d", Adapter: "openai", Models: []config.Model{{ID: "deepseek-v4-pro"}}}, config.Source{Name: "zen", BaseURL: "http://z", Adapter: "openai", Models: []config.Model{{ID: "deepseek-v4-flash-free"}}}, ) scope := []config.ModelScope{{Model: "deepseek-v4-flash-free"}, {Model: "gpt-5.6-sol"}} for _, prefixed := range []string{"zen:deepseek-v4-flash-free", "zen/deepseek-v4-flash-free"} { if !g.hasScopeModel(scope, prefixed) { t.Errorf("hasScopeModel(scope, %q)=false, want true (prefix must be stripped)", prefixed) } } for _, bare := range []string{"deepseek-v4-flash-free", "gpt-5.6-sol"} { if !g.hasScopeModel(scope, bare) { t.Errorf("hasScopeModel(scope, %q)=false, want true", bare) } } if g.hasScopeModel(scope, "deepseek-v4-pro") { t.Error("hasScopeModel returned true for model outside scope") } if g.hasScopeModel(scope, "deepseek-v4-flash-free-extra") { t.Error("hasScopeModel returned true for unrelated model") } if !g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:anything") { t.Error("AUTO scope should allow any prefixed model") } } // TestStatusRecorderFlusher pins the SSE-critical contract: the access-log // wrapper must implement http.Flusher, otherwise the streaming chat handlers' // w.(http.Flusher) assertion yields nil and chunks are never flushed until // the response ends (regression for the buffered-SSE bug). func TestStatusRecorderFlusher(t *testing.T) { var inner *flushRecorder rr := httptest.NewRecorder() inner = &flushRecorder{ResponseWriter: rr} sr := &statusRecorder{ResponseWriter: inner} fl, ok := any(sr).(http.Flusher) if !ok { t.Fatal("statusRecorder does not implement http.Flusher — SSE streaming is broken") } if inner.flushed { t.Fatal("Flush called before Flush()") } fl.Flush() if !inner.flushed { t.Fatal("statusRecorder.Flush did not forward to the underlying writer") } } // flushRecorder records whether Flush was forwarded. type flushRecorder struct { http.ResponseWriter flushed bool } func (f *flushRecorder) Flush() { f.flushed = true }