mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 17:07:59 +00:00
7fb8f96b82
A key with scope=[AUTO] could previously: 1. request ANY concrete model id directly (hasScopeModel/checkModelScope treated AUTO as a wildcard) 2. see the full 56-model list on /v1/models (intersectModels considered AUTO as grant-everything) AUTO now only authorizes the AUTO routing mode. Direct requests to a specific model require an explicit scope entry. Also carries agentrouter.lua WAF fingerprint headers (Origin/Referer/ X-Requested-With) already staged on this branch. Tests: TestHasScopeModelWithSourcePrefix updated; full suite green.
1021 lines
36 KiB
Go
1021 lines
36 KiB
Go
package gateway
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"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:deepseek-v4-flash-free") {
|
|
t.Error("AUTO scope must NOT grant specific model ids even with prefix")
|
|
}
|
|
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "deepseek-v4-pro") {
|
|
t.Error("AUTO scope must NOT grant specific model ids")
|
|
}
|
|
}
|
|
|
|
// 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 }
|
|
|
|
// TestDirectStreamFailoverAuditSource: a direct streaming request whose first
|
|
// candidate hard-fails must be served by the fallback, and the audit record's
|
|
// Source must name the source that actually served the stream (not just the
|
|
// first candidate).
|
|
func TestDirectStreamFailoverAuditSource(t *testing.T) {
|
|
aUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(500)
|
|
fmt.Fprint(w, `{"error":"boom"}`)
|
|
}))
|
|
defer aUp.Close()
|
|
bUp := mockUpstream()
|
|
defer bUp.Close()
|
|
g := newTestGateway(t,
|
|
// Both sources expose the same model id so the direct scheduler has a
|
|
// real fallback candidate after a hard-fails; b is the healthy one.
|
|
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
|
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
|
)
|
|
body := doReq(t, g, "POST", "/v1/chat/completions",
|
|
`{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
|
if body.Code != 200 {
|
|
t.Fatalf("status=%d body=%s", body.Code, body.Body.String())
|
|
}
|
|
sb := body.Body.String()
|
|
if !strings.Contains(sb, `"finish_reason":"stop"`) || !strings.Contains(sb, "[DONE]") {
|
|
t.Fatalf("stream incomplete: %s", sb)
|
|
}
|
|
recs := g.stats.AuditRecords(0, 0, "")
|
|
if len(recs) == 0 {
|
|
t.Fatal("no audit records")
|
|
}
|
|
last := recs[len(recs)-1]
|
|
if last.Source != "b" || !last.OK {
|
|
t.Fatalf("audit rec source=%q ok=%v, want source=b ok=true (actual serving source)", last.Source, last.OK)
|
|
}
|
|
}
|
|
|
|
// TestStatusExposesAdapterPools checks that the elastic Lua pool sizing is
|
|
// visible to admins (and only to them) via /api/status, so the grow/shrink
|
|
// algorithm can be observed instead of inferred.
|
|
func TestStatusExposesAdapterPools(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
rr := doReq(t, g, http.MethodGet, "/api/status", "")
|
|
if rr.Code != 200 {
|
|
t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
|
|
}
|
|
var body map[string]interface{}
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
|
|
t.Fatalf("decode: %v", err)
|
|
}
|
|
pools, ok := body["adapter_pools"].([]interface{})
|
|
if !ok {
|
|
t.Fatalf("adapter_pools missing or wrong type: %T", body["adapter_pools"])
|
|
}
|
|
if len(pools) == 0 {
|
|
t.Fatal("adapter_pools must list the loaded adapters")
|
|
}
|
|
first, _ := pools[0].(map[string]interface{})
|
|
for _, field := range []string{"name", "created", "idle", "in_use", "max", "grow_step", "shrink_step"} {
|
|
if _, ok := first[field]; !ok {
|
|
t.Errorf("pool entry missing %q: %v", field, first)
|
|
}
|
|
}
|
|
// idle gateway: nothing booted eagerly
|
|
if c, _ := first["created"].(float64); c != 0 {
|
|
t.Errorf("adapter pool booted eagerly, created=%v", first["created"])
|
|
}
|
|
}
|
|
|
|
// TestStatsRecordsAPIPaging exercises the on-demand records endpoint the
|
|
// dashboard scrolls with: a bounded first page plus a cursor that walks back
|
|
// through the audit file.
|
|
func TestStatsRecordsAPIPaging(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
for i := 0; i < 250; i++ {
|
|
g.stats.Record(Req{
|
|
Time: time.Now().UnixMilli() - int64(250-i)*1000,
|
|
Key: keyID("sk-test"), Type: "chat", Model: "m", Source: "s",
|
|
Prompt: 1, Compl: 1, OK: true, Status: 200,
|
|
})
|
|
}
|
|
|
|
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=50", "")
|
|
if rr.Code != 200 {
|
|
t.Fatalf("status = %d body=%s", rr.Code, rr.Body.String())
|
|
}
|
|
var page struct {
|
|
Records []Req `json:"records"`
|
|
NextCursor string `json:"next_cursor"`
|
|
HasMore bool `json:"has_more"`
|
|
}
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
|
|
t.Fatalf("decode: %v", err)
|
|
}
|
|
if len(page.Records) != 50 {
|
|
t.Fatalf("first page = %d records, want 50", len(page.Records))
|
|
}
|
|
if !page.HasMore || page.NextCursor == "" {
|
|
t.Fatal("a long history must advertise more pages plus a cursor")
|
|
}
|
|
// newest first
|
|
for i := 1; i < len(page.Records); i++ {
|
|
if page.Records[i-1].Time < page.Records[i].Time {
|
|
t.Fatalf("page not newest-first at %d", i)
|
|
}
|
|
}
|
|
|
|
// follow the cursor: the next page must continue strictly older
|
|
oldest := page.Records[len(page.Records)-1].Time
|
|
rr2 := doReq(t, g, http.MethodGet,
|
|
"/api/stats/records?limit=50&before="+url.QueryEscape(page.NextCursor), "")
|
|
if rr2.Code != 200 {
|
|
t.Fatalf("page 2 status = %d", rr2.Code)
|
|
}
|
|
var page2 struct {
|
|
Records []Req `json:"records"`
|
|
}
|
|
if err := json.Unmarshal(rr2.Body.Bytes(), &page2); err != nil {
|
|
t.Fatalf("decode page 2: %v", err)
|
|
}
|
|
if len(page2.Records) == 0 {
|
|
t.Fatal("cursor page must return records")
|
|
}
|
|
if page2.Records[0].Time > oldest {
|
|
t.Fatalf("page 2 starts at %d, must continue below %d", page2.Records[0].Time, oldest)
|
|
}
|
|
}
|
|
|
|
// TestStatsSnapshotShipsOneScreen: a dashboard poll must not serialize the whole
|
|
// ring, otherwise the "load the first screen only" contract is broken on the
|
|
// wire even if the UI pages.
|
|
func TestStatsSnapshotShipsOneScreen(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
for i := 0; i < 400; i++ {
|
|
g.stats.Record(Req{
|
|
Time: time.Now().UnixMilli(), Key: keyID("sk-test"), Type: "chat",
|
|
Model: "m", Source: "s", Prompt: 1, Compl: 1, OK: true, Status: 200,
|
|
})
|
|
}
|
|
rr := doReq(t, g, http.MethodGet, "/api/stats", "")
|
|
var snap struct {
|
|
Records []Req `json:"records"`
|
|
Total Stat `json:"total"`
|
|
}
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &snap); err != nil {
|
|
t.Fatalf("decode: %v", err)
|
|
}
|
|
if len(snap.Records) > firstScreenRecords {
|
|
t.Fatalf("snapshot shipped %d records, want <= %d", len(snap.Records), firstScreenRecords)
|
|
}
|
|
// aggregates still cover every request
|
|
if snap.Total.Reqs != 400 {
|
|
t.Fatalf("total reqs = %d, want 400 (aggregates must not be paged away)", snap.Total.Reqs)
|
|
}
|
|
}
|
|
|
|
// TestStatsCsvExportStreams checks the export still emits every row in the
|
|
// window now that it is written straight from the audit walk.
|
|
func TestStatsCsvExportStreams(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
base := time.Now().UnixMilli()
|
|
for i := 0; i < 120; i++ {
|
|
g.stats.Record(Req{
|
|
Time: base + int64(i)*1000, Key: keyID("sk-test"), Type: "chat",
|
|
Model: "m", Source: "s", Prompt: 2, Compl: 3, OK: true, Status: 200,
|
|
})
|
|
}
|
|
rr := doReq(t, g, http.MethodGet,
|
|
"/api/stats?export=csv&from="+strconv.FormatInt(base, 10)+
|
|
"&to="+strconv.FormatInt(base+120*1000, 10), "")
|
|
if rr.Code != 200 {
|
|
t.Fatalf("export status = %d", rr.Code)
|
|
}
|
|
lines := strings.Split(strings.TrimSpace(rr.Body.String()), "\n")
|
|
if len(lines) != 121 { // header + 120 rows
|
|
t.Fatalf("csv had %d lines, want 121 (header + 120 records)", len(lines))
|
|
}
|
|
if !strings.HasPrefix(lines[0], "time,key,key_name") {
|
|
t.Fatalf("unexpected csv header: %q", lines[0])
|
|
}
|
|
}
|
|
|
|
// TestStatsRecordsAPIScopedToOwnKey: a non-admin key must only page its own
|
|
// records even when it asks for someone else's.
|
|
func TestStatsRecordsAPIScopedToOwnKey(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
mine := keyID("sk-test")
|
|
for i := 0; i < 5; i++ {
|
|
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: mine, Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
|
|
}
|
|
for i := 0; i < 5; i++ {
|
|
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: "othermask", Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
|
|
}
|
|
// sk-test is an admin in the test gateway, so it legitimately sees all keys;
|
|
// assert the explicit filter path instead.
|
|
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=100&key=othermask", "")
|
|
var page struct {
|
|
Records []Req `json:"records"`
|
|
}
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
|
|
t.Fatalf("decode: %v", err)
|
|
}
|
|
for _, r := range page.Records {
|
|
if r.Key != "othermask" {
|
|
t.Fatalf("key filter leaked %q", r.Key)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestKeyDeleteRoundTrip is the end-to-end version of the bug users hit:
|
|
// deleting a gateway key from the admin UI. It also asserts the GET-on-a-key
|
|
// error names the right verb, since that message is what a mis-methoded client
|
|
// actually sees.
|
|
func TestKeyDeleteRoundTrip(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
|
|
rr := doReq(t, g, http.MethodPost, "/api/keys", `{"name":"doomed","role":"user"}`)
|
|
if rr.Code != 200 {
|
|
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
|
}
|
|
var created struct {
|
|
Key struct {
|
|
Key string `json:"key"`
|
|
Name string `json:"name"`
|
|
} `json:"key"`
|
|
}
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil {
|
|
t.Fatalf("decode create: %v", err)
|
|
}
|
|
if created.Key.Key == "" {
|
|
t.Fatalf("no key returned: %s", rr.Body.String())
|
|
}
|
|
|
|
// DELETE removes it
|
|
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
|
|
if rr.Code != 200 {
|
|
t.Fatalf("delete: %d %s", rr.Code, rr.Body.String())
|
|
}
|
|
for _, k := range g.core.ListKeys() {
|
|
if k.Key == created.Key.Key {
|
|
t.Fatal("key still present after DELETE")
|
|
}
|
|
}
|
|
|
|
// deleting again is a clean 404, not a 500
|
|
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
|
|
if rr.Code != http.StatusNotFound {
|
|
t.Fatalf("second delete: %d, want 404", rr.Code)
|
|
}
|
|
}
|
|
|
|
// TestKeyGetOnSpecificKeyNamesTheVerb: a GET on /api/keys/{key} is what a
|
|
// client that forgot to set method:"DELETE" ends up sending. The error must say
|
|
// which verb to use instead of just "use GET /api/keys", which reads as if the
|
|
// caller did nothing wrong.
|
|
func TestKeyGetOnSpecificKeyNamesTheVerb(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
rr := doReq(t, g, http.MethodGet, "/api/keys/sk-whatever", "")
|
|
if rr.Code != http.StatusNotFound {
|
|
t.Fatalf("status = %d, want 404", rr.Code)
|
|
}
|
|
body := rr.Body.String()
|
|
for _, want := range []string{"DELETE /api/keys/", "PUT /api/keys/"} {
|
|
if !strings.Contains(body, want) {
|
|
t.Errorf("error message should mention %q, got: %s", want, body)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestAdapterGetOnSpecificAdapterNamesDelete is the same guard for the adapters
|
|
// endpoint, whose delAdapter helper had the identical missing-method bug.
|
|
func TestAdapterGetOnSpecificAdapterNamesDelete(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
rr := doReq(t, g, http.MethodGet, "/api/adapters/openai", "")
|
|
if rr.Code != http.StatusNotFound {
|
|
t.Fatalf("status = %d, want 404", rr.Code)
|
|
}
|
|
if !strings.Contains(rr.Body.String(), "DELETE /api/adapters/openai") {
|
|
t.Errorf("error should name the DELETE verb, got: %s", rr.Body.String())
|
|
}
|
|
}
|
|
|
|
// TestCannotDeleteOwnKey keeps the admin from locking themselves out.
|
|
func TestCannotDeleteOwnKey(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
rr := doReq(t, g, http.MethodDelete, "/api/keys/sk-test", "")
|
|
if rr.Code != http.StatusBadRequest {
|
|
t.Fatalf("status = %d, want 400", rr.Code)
|
|
}
|
|
if !strings.Contains(rr.Body.String(), "logged in with") {
|
|
t.Errorf("unexpected message: %s", rr.Body.String())
|
|
}
|
|
}
|
|
|
|
// TestGatewayRingStaysSmall pins the resident record ring at the small default.
|
|
// The gateway used to construct Stats with a 10000-record ring, which after the
|
|
// full-history aggregate scan meant 10000 resident Req structs — putting back the
|
|
// ~25 MB the on-demand paging work removed. Aggregates cover all of history;
|
|
// the ring only needs the first screen plus the 5-minute status windows.
|
|
func TestGatewayRingStaysSmall(t *testing.T) {
|
|
g := newTestGateway(t)
|
|
g.stats.mu.Lock()
|
|
max := g.stats.maxRecs
|
|
g.stats.mu.Unlock()
|
|
if max != defaultRingSize {
|
|
t.Fatalf("gateway record ring = %d, want %d (a big ring reintroduces the resident-memory cost)",
|
|
max, defaultRingSize)
|
|
}
|
|
}
|