LovroG05
8bb62a1c20
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Any model can name another model as its backup. The gateway tries the backup when the model is offline, rate limited, or its upstream fails before the client has received anything. The response reports the requested model, so the client sees no switch. The request log records each failed attempt and the address that served the request, and the cost uses the backup's prices.
670 lines
24 KiB
Go
670 lines
24 KiB
Go
package api
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/catalog"
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/providermgr"
|
|
_ "git.govekar.net/LovroG05/localmodelroutex/internal/providers/all"
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/ratelimit"
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
|
)
|
|
|
|
// The proxy handlers are exercised end to end against a fake upstream: a real
|
|
// OpenRouter provider built by the provider manager, a real catalog cache, a
|
|
// real rate limiter, and a fake request log sink in place of Postgres.
|
|
|
|
type fakeCatalogStore struct{ rows []store.CatalogRow }
|
|
|
|
func (f *fakeCatalogStore) ListCatalog(context.Context) ([]store.CatalogRow, error) {
|
|
return f.rows, nil
|
|
}
|
|
func (f *fakeCatalogStore) GetSetting(context.Context, string, any) (bool, error) {
|
|
return false, nil
|
|
}
|
|
|
|
type fakeProviderStore struct{ rows []store.ProviderRow }
|
|
|
|
func (f *fakeProviderStore) ListProviders(context.Context, bool) ([]store.ProviderRow, error) {
|
|
return f.rows, nil
|
|
}
|
|
|
|
// fakeLogs records log rows and signals on done once the content row lands,
|
|
// which is the last write logRequest performs.
|
|
type fakeLogs struct {
|
|
mu sync.Mutex
|
|
rows []store.RequestLog
|
|
messages []string
|
|
contents []string
|
|
touched []uuid.UUID
|
|
done chan struct{}
|
|
}
|
|
|
|
func newFakeLogs() *fakeLogs { return &fakeLogs{done: make(chan struct{}, 16)} }
|
|
|
|
func (f *fakeLogs) InsertRequestLog(_ context.Context, e store.RequestLog) (uuid.UUID, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.rows = append(f.rows, e)
|
|
return uuid.New(), nil
|
|
}
|
|
|
|
func (f *fakeLogs) InsertLogContent(_ context.Context, _ uuid.UUID, messages json.RawMessage, completion string) error {
|
|
f.mu.Lock()
|
|
f.messages = append(f.messages, string(messages))
|
|
f.contents = append(f.contents, completion)
|
|
f.mu.Unlock()
|
|
f.done <- struct{}{}
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeLogs) TouchVirtualKey(_ context.Context, id uuid.UUID) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.touched = append(f.touched, id)
|
|
return nil
|
|
}
|
|
|
|
// waitN blocks until n requests have been logged and returns the rows in the
|
|
// order they were written.
|
|
func (f *fakeLogs) waitN(t *testing.T, n int) []store.RequestLog {
|
|
t.Helper()
|
|
for range n {
|
|
select {
|
|
case <-f.done:
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("request log was not written")
|
|
}
|
|
}
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return append([]store.RequestLog(nil), f.rows...)
|
|
}
|
|
|
|
// wait blocks until one request has been fully logged.
|
|
func (f *fakeLogs) wait(t *testing.T) store.RequestLog {
|
|
t.Helper()
|
|
rows := f.waitN(t, 1)
|
|
return rows[len(rows)-1]
|
|
}
|
|
|
|
// upstream is a fake OpenRouter that records the last body it received.
|
|
type upstream struct {
|
|
mu sync.Mutex
|
|
calls int
|
|
body map[string]any
|
|
// handler decides the reply; set per test.
|
|
handler func(w http.ResponseWriter, r *http.Request, body map[string]any)
|
|
}
|
|
|
|
func (u *upstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
var body map[string]any
|
|
b, _ := io.ReadAll(r.Body)
|
|
_ = json.Unmarshal(b, &body)
|
|
u.mu.Lock()
|
|
u.calls++
|
|
u.body = body
|
|
u.mu.Unlock()
|
|
u.handler(w, r, body)
|
|
}
|
|
|
|
func (u *upstream) last() (int, map[string]any) {
|
|
u.mu.Lock()
|
|
defer u.mu.Unlock()
|
|
return u.calls, u.body
|
|
}
|
|
|
|
const testPNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
|
|
|
|
func imageReply(w http.ResponseWriter, _ *http.Request, _ map[string]any) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"id":"gen","choices":[{"message":{"role":"assistant","content":"done","images":[{"type":"image_url","image_url":{"url":"data:image/png;base64,`+testPNG+`"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":1290,"total_tokens":1300,"cost":0.03}}`)
|
|
}
|
|
|
|
func chatReply(w http.ResponseWriter, _ *http.Request, body map[string]any) {
|
|
if body["stream"] == true {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
_, _ = io.WriteString(w, "data: {\"id\":\"c\",\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
|
|
_, _ = io.WriteString(w, "data: {\"id\":\"c\",\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n")
|
|
_, _ = io.WriteString(w, "data: {\"id\":\"c\",\"choices\":[],\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":50,\"total_tokens\":150}}\n\n")
|
|
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"id":"c","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":100,"completion_tokens":50,"total_tokens":150}}`)
|
|
}
|
|
|
|
type testEnv struct {
|
|
api *API
|
|
logs *fakeLogs
|
|
upstream *upstream
|
|
}
|
|
|
|
// newTestEnv wires an API with one OpenRouter provider ("or") serving an
|
|
// image model and a text model.
|
|
func newTestEnv(t *testing.T, handler func(http.ResponseWriter, *http.Request, map[string]any), tweak func(rows []store.CatalogRow)) *testEnv {
|
|
t.Helper()
|
|
up := &upstream{handler: handler}
|
|
srv := httptest.NewServer(up)
|
|
t.Cleanup(srv.Close)
|
|
|
|
providerID := uuid.New()
|
|
mgr := providermgr.New(&fakeProviderStore{rows: []store.ProviderRow{
|
|
{ID: providerID, Name: "or", Type: "openrouter", BaseURL: srv.URL, Enabled: true},
|
|
}}, nil)
|
|
if err := mgr.Reload(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
rows := []store.CatalogRow{
|
|
{
|
|
ID: uuid.New(), ProviderID: providerID, ProviderName: "or", ProviderType: "openrouter",
|
|
Slug: "google/gemini-2.5-flash-image", DisplayName: "Nano Banana",
|
|
Capabilities: []string{"images", "chat"}, Enabled: true, Available: true,
|
|
InputPricePerMTok: 0.3, OutputPricePerMTok: 2.5,
|
|
},
|
|
{
|
|
ID: uuid.New(), ProviderID: providerID, ProviderName: "or", ProviderType: "openrouter",
|
|
Slug: "openai/gpt-4o", DisplayName: "GPT-4o",
|
|
Capabilities: []string{"chat"}, Enabled: true, Available: true,
|
|
InputPricePerMTok: 0.3, OutputPricePerMTok: 2.5,
|
|
},
|
|
}
|
|
if tweak != nil {
|
|
tweak(rows)
|
|
}
|
|
cache := catalog.NewCache(&fakeCatalogStore{rows: rows})
|
|
if err := cache.Reload(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
logs := newFakeLogs()
|
|
return &testEnv{
|
|
api: New(Deps{Cache: cache, Manager: mgr, Limiter: ratelimit.New(), Logs: logs}),
|
|
logs: logs,
|
|
upstream: up,
|
|
}
|
|
}
|
|
|
|
func allowAllKey() *store.VirtualKey {
|
|
return &store.VirtualKey{ID: uuid.New(), KeyPrefix: "sk-test", AllowAllModels: true}
|
|
}
|
|
|
|
// call invokes a proxy handler as the key middleware would have.
|
|
func call(h http.HandlerFunc, key *store.VirtualKey, method, body string) *httptest.ResponseRecorder {
|
|
req := httptest.NewRequest(method, "/v1/x", bytes.NewBufferString(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req = req.WithContext(context.WithValue(req.Context(), keyCtxKey, key))
|
|
rec := httptest.NewRecorder()
|
|
h(rec, req)
|
|
return rec
|
|
}
|
|
|
|
func TestImageGenerationsEndToEnd(t *testing.T) {
|
|
env := newTestEnv(t, imageReply, func(rows []store.CatalogRow) {
|
|
rows[0].ParamOverrides = []byte(`{"image_config":{"image_size":"2K"}}`)
|
|
})
|
|
key := allowAllKey()
|
|
rec := call(env.api.imageGenerations, key, http.MethodPost,
|
|
`{"model":"or/google/gemini-2.5-flash-image","prompt":"a red cube","size":"1792x1024","response_format":"b64_json"}`)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status = %d body=%s", rec.Code, rec.Body)
|
|
}
|
|
|
|
var out struct {
|
|
Data []struct {
|
|
B64 string `json:"b64_json"`
|
|
RevisedPrompt string `json:"revised_prompt"`
|
|
} `json:"data"`
|
|
Usage struct {
|
|
Cost float64 `json:"cost"`
|
|
} `json:"usage"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(out.Data) != 1 || out.Data[0].B64 != testPNG || out.Data[0].RevisedPrompt != "done" || out.Usage.Cost != 0.03 {
|
|
t.Errorf("body = %s", rec.Body)
|
|
}
|
|
|
|
// Upstream saw the translated chat request with the slug and the override.
|
|
calls, body := env.upstream.last()
|
|
if calls != 1 || body["model"] != "google/gemini-2.5-flash-image" {
|
|
t.Errorf("upstream calls=%d model=%v", calls, body["model"])
|
|
}
|
|
cfg, _ := body["image_config"].(map[string]any)
|
|
if cfg["aspect_ratio"] != "16:9" || cfg["image_size"] != "2K" {
|
|
t.Errorf("image_config = %v (override must merge with mapped size)", cfg)
|
|
}
|
|
if mods, _ := body["modalities"].([]any); len(mods) != 2 {
|
|
t.Errorf("modalities = %v", body["modalities"])
|
|
}
|
|
|
|
// Accounting: upstream cost wins, the prompt is logged, the image bytes are not.
|
|
row := env.logs.wait(t)
|
|
if row.ModelName != "or/google/gemini-2.5-flash-image" || row.ProviderName != "or" || row.StatusCode != 200 {
|
|
t.Errorf("log row = %+v", row)
|
|
}
|
|
if row.Cost != 0.03 || row.TotalTokens != 1300 || row.Streamed {
|
|
t.Errorf("accounting = cost %v tokens %d streamed %v", row.Cost, row.TotalTokens, row.Streamed)
|
|
}
|
|
if row.KeyPrefix != "sk-test" || row.VirtualKeyID == nil || *row.VirtualKeyID != key.ID {
|
|
t.Errorf("key attribution = %+v", row)
|
|
}
|
|
if !strings.Contains(env.logs.messages[0], "a red cube") {
|
|
t.Errorf("prompt not logged: %s", env.logs.messages[0])
|
|
}
|
|
if !strings.Contains(env.logs.contents[0], "1 image(s) generated") || strings.Contains(env.logs.contents[0], testPNG) {
|
|
t.Errorf("completion log = %q", env.logs.contents[0])
|
|
}
|
|
if len(env.logs.touched) != 1 || env.logs.touched[0] != key.ID {
|
|
t.Errorf("key last-used not touched: %v", env.logs.touched)
|
|
}
|
|
}
|
|
|
|
func TestImageGenerationsRejectsTextOnlyModel(t *testing.T) {
|
|
env := newTestEnv(t, imageReply, nil)
|
|
rec := call(env.api.imageGenerations, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","prompt":"x"}`)
|
|
if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), "does not generate images") {
|
|
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
if calls, _ := env.upstream.last(); calls != 0 {
|
|
t.Errorf("upstream must not be called, got %d calls", calls)
|
|
}
|
|
}
|
|
|
|
func TestImageGenerationsValidation(t *testing.T) {
|
|
env := newTestEnv(t, imageReply, nil)
|
|
cases := map[string]struct {
|
|
body string
|
|
want int
|
|
msg string
|
|
}{
|
|
"missing model": {`{"prompt":"x"}`, http.StatusBadRequest, "missing 'model'"},
|
|
"missing prompt": {`{"model":"or/google/gemini-2.5-flash-image"}`, http.StatusBadRequest, "missing 'prompt'"},
|
|
"too many": {`{"model":"or/google/gemini-2.5-flash-image","prompt":"x","n":11}`, http.StatusBadRequest, "at most 10"},
|
|
"unknown model": {`{"model":"or/nope","prompt":"x"}`, http.StatusNotFound, ""},
|
|
"bad address": {`{"model":"nope","prompt":"x"}`, http.StatusNotFound, "invalid model address"},
|
|
"bad json": {`{`, http.StatusBadRequest, "decode"},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
rec := call(env.api.imageGenerations, allowAllKey(), http.MethodPost, tc.body)
|
|
if rec.Code != tc.want || !strings.Contains(rec.Body.String(), tc.msg) {
|
|
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
})
|
|
}
|
|
if calls, _ := env.upstream.last(); calls != 0 {
|
|
t.Errorf("upstream must not be called, got %d calls", calls)
|
|
}
|
|
}
|
|
|
|
func TestImageGenerationsAllowlist(t *testing.T) {
|
|
env := newTestEnv(t, imageReply, nil)
|
|
key := &store.VirtualKey{ID: uuid.New(), KeyPrefix: "sk-x", Rules: []string{"or/openai/*"}}
|
|
rec := call(env.api.imageGenerations, key, http.MethodPost,
|
|
`{"model":"or/google/gemini-2.5-flash-image","prompt":"x"}`)
|
|
if rec.Code != http.StatusForbidden {
|
|
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
}
|
|
|
|
func TestImageGenerationsFanOutCountsAgainstRPM(t *testing.T) {
|
|
env := newTestEnv(t, imageReply, func(rows []store.CatalogRow) {
|
|
two := 2
|
|
rows[0].RPMLimit = &two
|
|
})
|
|
body := `{"model":"or/google/gemini-2.5-flash-image","prompt":"x","n":3}`
|
|
if rec := call(env.api.imageGenerations, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusTooManyRequests {
|
|
t.Errorf("n=3 against limit 2: status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
body = `{"model":"or/google/gemini-2.5-flash-image","prompt":"x","n":2}`
|
|
if rec := call(env.api.imageGenerations, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusOK {
|
|
t.Errorf("n=2 against limit 2: status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
env.logs.wait(t)
|
|
if calls, _ := env.upstream.last(); calls != 2 {
|
|
t.Errorf("upstream calls = %d, want 2", calls)
|
|
}
|
|
}
|
|
|
|
func TestImageGenerationsRelaysUpstreamError(t *testing.T) {
|
|
env := newTestEnv(t, func(w http.ResponseWriter, _ *http.Request, _ map[string]any) {
|
|
w.WriteHeader(http.StatusPaymentRequired)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"Insufficient credits"}}`)
|
|
}, nil)
|
|
rec := call(env.api.imageGenerations, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/google/gemini-2.5-flash-image","prompt":"x"}`)
|
|
if rec.Code != http.StatusPaymentRequired {
|
|
t.Errorf("status=%d", rec.Code)
|
|
}
|
|
row := env.logs.wait(t)
|
|
if row.StatusCode != http.StatusPaymentRequired || row.Error != "Insufficient credits" {
|
|
t.Errorf("log row = %+v", row)
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsEndToEnd(t *testing.T) {
|
|
env := newTestEnv(t, chatReply, nil)
|
|
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
var out map[string]any
|
|
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
|
if _, ok := out["created"]; !ok {
|
|
t.Errorf("created not injected: %s", rec.Body)
|
|
}
|
|
if _, body := env.upstream.last(); body["model"] != "openai/gpt-4o" {
|
|
t.Errorf("upstream model = %v", body["model"])
|
|
}
|
|
row := env.logs.wait(t)
|
|
// No upstream cost: 100 prompt tokens at $0.3/M plus 50 completion at $2.5/M.
|
|
want := 100.0/1e6*0.3 + 50.0/1e6*2.5
|
|
if row.PromptTokens != 100 || row.CompletionTokens != 50 || row.Cost < want-1e-12 || row.Cost > want+1e-12 {
|
|
t.Errorf("accounting = %+v, want cost %v", row, want)
|
|
}
|
|
if env.logs.contents[0] != "hello" {
|
|
t.Errorf("completion log = %q", env.logs.contents[0])
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsStream(t *testing.T) {
|
|
env := newTestEnv(t, chatReply, nil)
|
|
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","messages":[{"role":"user","content":"hi"}],"stream":true}`)
|
|
if rec.Code != http.StatusOK || rec.Header().Get("Content-Type") != "text/event-stream" {
|
|
t.Fatalf("status=%d type=%q", rec.Code, rec.Header().Get("Content-Type"))
|
|
}
|
|
body := rec.Body.String()
|
|
if strings.Count(body, "data: ") != 4 || !strings.HasSuffix(body, "data: [DONE]\n\n") {
|
|
t.Errorf("stream body = %q", body)
|
|
}
|
|
if _, up := env.upstream.last(); up["stream"] != true {
|
|
t.Errorf("upstream must be asked to stream: %v", up)
|
|
}
|
|
row := env.logs.wait(t)
|
|
if !row.Streamed || row.TTFTMs == nil || row.TotalTokens != 150 {
|
|
t.Errorf("log row = %+v", row)
|
|
}
|
|
if env.logs.contents[0] != "ab" {
|
|
t.Errorf("accumulated completion = %q", env.logs.contents[0])
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsRateLimited(t *testing.T) {
|
|
env := newTestEnv(t, chatReply, func(rows []store.CatalogRow) {
|
|
one := 1
|
|
rows[1].RPMLimit = &one
|
|
})
|
|
body := `{"model":"or/openai/gpt-4o","messages":[]}`
|
|
if rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusOK {
|
|
t.Fatalf("first call status=%d", rec.Code)
|
|
}
|
|
env.logs.wait(t)
|
|
if rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusTooManyRequests {
|
|
t.Errorf("second call status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
}
|
|
|
|
func TestEmbeddingsEndToEnd(t *testing.T) {
|
|
env := newTestEnv(t, func(w http.ResponseWriter, r *http.Request, _ map[string]any) {
|
|
if r.URL.Path != "/embeddings" {
|
|
w.WriteHeader(http.StatusNotFound)
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, `{"object":"list","data":[{"embedding":[0.1]}],"usage":{"prompt_tokens":4,"total_tokens":4}}`)
|
|
}, nil)
|
|
rec := call(env.api.embeddings, allowAllKey(), http.MethodPost, `{"model":"or/openai/gpt-4o","input":"x"}`)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
if row := env.logs.wait(t); row.TotalTokens != 4 {
|
|
t.Errorf("log row = %+v", row)
|
|
}
|
|
}
|
|
|
|
func TestListModelsHonoursAllowlistAndExposesCapabilities(t *testing.T) {
|
|
env := newTestEnv(t, chatReply, nil)
|
|
key := &store.VirtualKey{ID: uuid.New(), Rules: []string{"or/google/*"}}
|
|
rec := call(env.api.listModels, key, http.MethodGet, "")
|
|
var out struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
Capabilities []string `json:"capabilities"`
|
|
DisplayName string `json:"display_name"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(out.Data) != 1 || out.Data[0].ID != "or/google/gemini-2.5-flash-image" || out.Data[0].DisplayName != "Nano Banana" {
|
|
t.Fatalf("data = %+v", out.Data)
|
|
}
|
|
if len(out.Data[0].Capabilities) != 2 {
|
|
t.Errorf("capabilities = %v", out.Data[0].Capabilities)
|
|
}
|
|
}
|
|
|
|
func TestAllowedPatterns(t *testing.T) {
|
|
cases := []struct {
|
|
rules []string
|
|
addr string
|
|
want bool
|
|
}{
|
|
{[]string{"*"}, "or/x", true},
|
|
{[]string{"or/*"}, "or/google/x", true},
|
|
{[]string{"or/*"}, "local/x", false},
|
|
{[]string{"or/google/x"}, "or/google/x", true},
|
|
{[]string{"or/google/x"}, "or/google/y", false},
|
|
{nil, "or/x", false},
|
|
}
|
|
for _, tc := range cases {
|
|
k := &store.VirtualKey{Rules: tc.rules}
|
|
if got := allowed(k, tc.addr); got != tc.want {
|
|
t.Errorf("allowed(%v, %q) = %v", tc.rules, tc.addr, got)
|
|
}
|
|
}
|
|
if !allowed(&store.VirtualKey{AllowAllModels: true}, "anything/at/all") {
|
|
t.Error("allow-all key must pass")
|
|
}
|
|
}
|
|
|
|
// failPrimary answers gpt-4o with status and lets every other model through
|
|
// to chatReply.
|
|
func failPrimary(status int) func(http.ResponseWriter, *http.Request, map[string]any) {
|
|
return func(w http.ResponseWriter, r *http.Request, body map[string]any) {
|
|
if body["model"] == "openai/gpt-4o" {
|
|
w.WriteHeader(status)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"primary down"}}`)
|
|
return
|
|
}
|
|
chatReply(w, r, body)
|
|
}
|
|
}
|
|
|
|
// withBackup makes the image model the backup of gpt-4o, at a different price.
|
|
func withBackup(rows []store.CatalogRow) {
|
|
id := rows[0].ID
|
|
rows[1].FallbackModelID = &id
|
|
rows[0].InputPricePerMTok = 1
|
|
rows[0].OutputPricePerMTok = 2
|
|
}
|
|
|
|
// rowsByStatus indexes log rows by status code.
|
|
func rowsByStatus(rows []store.RequestLog) map[int]store.RequestLog {
|
|
out := map[int]store.RequestLog{}
|
|
for _, r := range rows {
|
|
out[r.StatusCode] = r
|
|
}
|
|
return out
|
|
}
|
|
|
|
func TestChatCompletionsFallsBackOnUpstreamError(t *testing.T) {
|
|
env := newTestEnv(t, failPrimary(http.StatusServiceUnavailable), withBackup)
|
|
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
// The client sees the model it asked for.
|
|
var out map[string]any
|
|
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
|
if out["model"] != "openai/gpt-4o" {
|
|
t.Errorf("reported model = %v, want the requested one", out["model"])
|
|
}
|
|
if calls, body := env.upstream.last(); calls != 2 || body["model"] != "google/gemini-2.5-flash-image" {
|
|
t.Errorf("upstream calls=%d last model=%v", calls, body["model"])
|
|
}
|
|
|
|
// The log sees both attempts and the backup's cost.
|
|
rows := rowsByStatus(env.logs.waitN(t, 2))
|
|
failed, ok := rows[http.StatusServiceUnavailable]
|
|
if !ok || !strings.Contains(failed.Error, "primary down") || !strings.Contains(failed.Error, "switched to backup or/google/gemini-2.5-flash-image") {
|
|
t.Errorf("failed attempt row = %+v", failed)
|
|
}
|
|
if failed.ServedModel != "" || failed.ModelName != "or/openai/gpt-4o" {
|
|
t.Errorf("failed attempt must be logged against the requested model: %+v", failed)
|
|
}
|
|
served := rows[http.StatusOK]
|
|
if served.ModelName != "or/openai/gpt-4o" || served.ServedModel != "or/google/gemini-2.5-flash-image" {
|
|
t.Errorf("served row = %+v", served)
|
|
}
|
|
want := 100.0/1e6*1 + 50.0/1e6*2
|
|
if served.Cost < want-1e-12 || served.Cost > want+1e-12 {
|
|
t.Errorf("cost = %v, want the backup's price %v", served.Cost, want)
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsStreamFallsBack(t *testing.T) {
|
|
env := newTestEnv(t, failPrimary(http.StatusBadGateway), withBackup)
|
|
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","messages":[],"stream":true}`)
|
|
if rec.Code != http.StatusOK || rec.Header().Get("Content-Type") != "text/event-stream" {
|
|
t.Fatalf("status=%d type=%q body=%s", rec.Code, rec.Header().Get("Content-Type"), rec.Body)
|
|
}
|
|
body := rec.Body.String()
|
|
if strings.Count(body, "data: ") != 4 || strings.Count(body, `"model":"openai/gpt-4o"`) != 3 {
|
|
t.Errorf("stream body = %q", body)
|
|
}
|
|
rows := rowsByStatus(env.logs.waitN(t, 2))
|
|
if rows[http.StatusOK].ServedModel != "or/google/gemini-2.5-flash-image" || !rows[http.StatusOK].Streamed {
|
|
t.Errorf("served row = %+v", rows[http.StatusOK])
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsUnavailableModelServedByBackup(t *testing.T) {
|
|
env := newTestEnv(t, chatReply, func(rows []store.CatalogRow) {
|
|
withBackup(rows)
|
|
rows[1].Available = false
|
|
})
|
|
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","messages":[]}`)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
var out map[string]any
|
|
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
|
if out["model"] != "openai/gpt-4o" {
|
|
t.Errorf("reported model = %v", out["model"])
|
|
}
|
|
if calls, _ := env.upstream.last(); calls != 1 {
|
|
t.Errorf("upstream calls = %d, want 1", calls)
|
|
}
|
|
row := env.logs.wait(t)
|
|
if row.ModelName != "or/openai/gpt-4o" || row.ServedModel != "or/google/gemini-2.5-flash-image" {
|
|
t.Errorf("log row = %+v", row)
|
|
}
|
|
|
|
// The down model stays listed, because clients can still call it.
|
|
list := call(env.api.listModels, allowAllKey(), http.MethodGet, "")
|
|
if !strings.Contains(list.Body.String(), `"or/openai/gpt-4o"`) {
|
|
t.Errorf("models list = %s", list.Body)
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsRateLimitOverflowsToBackup(t *testing.T) {
|
|
env := newTestEnv(t, chatReply, func(rows []store.CatalogRow) {
|
|
withBackup(rows)
|
|
one := 1
|
|
rows[1].RPMLimit = &one
|
|
})
|
|
body := `{"model":"or/openai/gpt-4o","messages":[]}`
|
|
if rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusOK {
|
|
t.Fatalf("first call status=%d", rec.Code)
|
|
}
|
|
if row := env.logs.wait(t); row.ServedModel != "" {
|
|
t.Errorf("first call must use the requested model: %+v", row)
|
|
}
|
|
if rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost, body); rec.Code != http.StatusOK {
|
|
t.Fatalf("second call status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
if row := env.logs.wait(t); row.ServedModel != "or/google/gemini-2.5-flash-image" {
|
|
t.Errorf("second call must overflow to the backup: %+v", row)
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsClientErrorIsNotRetried(t *testing.T) {
|
|
env := newTestEnv(t, failPrimary(http.StatusBadRequest), withBackup)
|
|
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","messages":[]}`)
|
|
if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), "primary down") {
|
|
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
if calls, _ := env.upstream.last(); calls != 1 {
|
|
t.Errorf("upstream calls = %d, want 1", calls)
|
|
}
|
|
if row := env.logs.wait(t); strings.Contains(row.Error, "switched") {
|
|
t.Errorf("log row = %+v", row)
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsLastFailureIsRelayed(t *testing.T) {
|
|
// Both the model and its backup are down: the client gets the last error.
|
|
env := newTestEnv(t, func(w http.ResponseWriter, _ *http.Request, _ map[string]any) {
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"all down"}}`)
|
|
}, withBackup)
|
|
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
|
`{"model":"or/openai/gpt-4o","messages":[]}`)
|
|
if rec.Code != http.StatusServiceUnavailable || !strings.Contains(rec.Body.String(), "all down") {
|
|
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
if calls, _ := env.upstream.last(); calls != 2 {
|
|
t.Errorf("upstream calls = %d, want 2", calls)
|
|
}
|
|
env.logs.waitN(t, 2)
|
|
}
|
|
|
|
func TestEmbeddingsFallBack(t *testing.T) {
|
|
env := newTestEnv(t, func(w http.ResponseWriter, _ *http.Request, body map[string]any) {
|
|
if body["model"] == "openai/gpt-4o" {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, `{"object":"list","model":"backup","data":[{"embedding":[0.1]}],"usage":{"prompt_tokens":4,"total_tokens":4}}`)
|
|
}, withBackup)
|
|
rec := call(env.api.embeddings, allowAllKey(), http.MethodPost, `{"model":"or/openai/gpt-4o","input":"x"}`)
|
|
if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), `"model":"openai/gpt-4o"`) {
|
|
t.Errorf("status=%d body=%s", rec.Code, rec.Body)
|
|
}
|
|
rows := rowsByStatus(env.logs.waitN(t, 2))
|
|
if rows[http.StatusOK].ServedModel != "or/google/gemini-2.5-flash-image" {
|
|
t.Errorf("served row = %+v", rows[http.StatusOK])
|
|
}
|
|
}
|