test: cover the proxy handlers and the OpenAI transport
Add a RequestLogSink interface so the handlers can be tested without Postgres. The tests run the image, chat, stream, embeddings and model routes end to end against a fake upstream.
This commit is contained in:
@@ -4,10 +4,12 @@ package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/google/uuid"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/catalog"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/crypto"
|
||||
@@ -43,9 +45,19 @@ func userFrom(ctx context.Context) *authedUser {
|
||||
return u
|
||||
}
|
||||
|
||||
// RequestLogSink receives the per-request accounting rows written by the
|
||||
// proxy handlers. *store.Store satisfies it; tests supply a fake.
|
||||
type RequestLogSink interface {
|
||||
InsertRequestLog(ctx context.Context, e store.RequestLog) (uuid.UUID, error)
|
||||
InsertLogContent(ctx context.Context, logID uuid.UUID, messages json.RawMessage, completion string) error
|
||||
TouchVirtualKey(ctx context.Context, id uuid.UUID) error
|
||||
}
|
||||
|
||||
// Deps are the dependencies the API handlers need.
|
||||
type Deps struct {
|
||||
Store *store.Store
|
||||
Store *store.Store
|
||||
// Logs receives request log rows. Defaults to Store when nil.
|
||||
Logs RequestLogSink
|
||||
Cache *catalog.Cache
|
||||
Manager *providermgr.Manager
|
||||
Cipher *crypto.Cipher
|
||||
@@ -74,7 +86,12 @@ type API struct {
|
||||
}
|
||||
|
||||
// New returns an API.
|
||||
func New(d Deps) *API { return &API{Deps: d} }
|
||||
func New(d Deps) *API {
|
||||
if d.Logs == nil && d.Store != nil {
|
||||
d.Logs = d.Store
|
||||
}
|
||||
return &API{Deps: d}
|
||||
}
|
||||
|
||||
// Mount registers the proxy and admin routes on r.
|
||||
func (a *API) Mount(r chi.Router) {
|
||||
|
||||
@@ -210,19 +210,20 @@ func (a *API) logRequest(r *http.Request, key *store.VirtualKey, entry *catalog.
|
||||
}
|
||||
}
|
||||
|
||||
go func() {
|
||||
// The log write must outlive the request, so it gets its own context.
|
||||
go func() { //nolint:gosec // G118: intentional, see above
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
id, err := a.Store.InsertRequestLog(ctx, rec)
|
||||
id, err := a.Logs.InsertRequestLog(ctx, rec)
|
||||
if err != nil {
|
||||
slog.Warn("failed to write request log", "err", err)
|
||||
return
|
||||
}
|
||||
if err := a.Store.InsertLogContent(ctx, id, messages, completion); err != nil {
|
||||
if err := a.Logs.InsertLogContent(ctx, id, messages, completion); err != nil {
|
||||
slog.Warn("failed to write request log content", "err", err)
|
||||
}
|
||||
if key != nil && key.ID != uuid.Nil {
|
||||
_ = a.Store.TouchVirtualKey(ctx, key.ID)
|
||||
_ = a.Logs.TouchVirtualKey(ctx, key.ID)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
475
backend/internal/api/proxy_test.go
Normal file
475
backend/internal/api/proxy_test.go
Normal file
@@ -0,0 +1,475 @@
|
||||
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
|
||||
}
|
||||
|
||||
// wait blocks until one request has been fully logged.
|
||||
func (f *fakeLogs) wait(t *testing.T) store.RequestLog {
|
||||
t.Helper()
|
||||
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 f.rows[len(f.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")
|
||||
}
|
||||
}
|
||||
212
backend/internal/providers/openaicompat/transport_test.go
Normal file
212
backend/internal/providers/openaicompat/transport_test.go
Normal file
@@ -0,0 +1,212 @@
|
||||
package openaicompat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/providers"
|
||||
)
|
||||
|
||||
type captured struct {
|
||||
path string
|
||||
headers http.Header
|
||||
body map[string]any
|
||||
}
|
||||
|
||||
// newTransport returns a Transport against a fake upstream and a pointer that
|
||||
// receives the last request the upstream saw.
|
||||
func newTransport(t *testing.T, status int, reply string, contentType string) (*Transport, *captured) {
|
||||
t.Helper()
|
||||
got := &captured{}
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
got.path = r.URL.Path
|
||||
got.headers = r.Header.Clone()
|
||||
got.body = nil
|
||||
if r.Body != nil {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
if len(b) > 0 {
|
||||
_ = json.Unmarshal(b, &got.body)
|
||||
}
|
||||
}
|
||||
if contentType != "" {
|
||||
w.Header().Set("Content-Type", contentType)
|
||||
}
|
||||
w.WriteHeader(status)
|
||||
_, _ = io.WriteString(w, reply)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return &Transport{
|
||||
Provider: "p",
|
||||
BaseURL: srv.URL + "/v1/", // trailing slash must be tolerated
|
||||
HTTP: srv.Client(),
|
||||
Authorize: BearerAuth("tok"),
|
||||
Headers: http.Header{"X-Title": []string{"LMR"}},
|
||||
IncludeUsage: true,
|
||||
IncludeCost: true,
|
||||
}, got
|
||||
}
|
||||
|
||||
func TestChatCompletionRelaysBodyAndUsage(t *testing.T) {
|
||||
tr, got := newTransport(t, http.StatusOK,
|
||||
`{"id":"x","choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5,"cost":0.001}}`,
|
||||
"application/json")
|
||||
req, _ := providers.ParseChatRequest([]byte(`{"model":"client-name","messages":[{"role":"user","content":"hey"}],"stream":true}`))
|
||||
req.Model = "upstream-name"
|
||||
|
||||
resp, err := tr.ChatCompletion(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.path != "/v1/chat/completions" {
|
||||
t.Errorf("path = %q", got.path)
|
||||
}
|
||||
if got.headers.Get("Authorization") != "Bearer tok" || got.headers.Get("X-Title") != "LMR" {
|
||||
t.Errorf("headers = %v", got.headers)
|
||||
}
|
||||
if got.body["model"] != "upstream-name" {
|
||||
t.Errorf("model = %v", got.body["model"])
|
||||
}
|
||||
if _, ok := got.body["stream"]; ok {
|
||||
t.Errorf("non-streaming call must drop stream=true")
|
||||
}
|
||||
if u, _ := got.body["usage"].(map[string]any); u["include"] != true {
|
||||
t.Errorf("usage.include missing: %v", got.body["usage"])
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK || resp.Usage.TotalTokens != 5 || resp.Usage.Cost == nil || *resp.Usage.Cost != 0.001 {
|
||||
t.Errorf("resp = %+v usage=%+v", resp.StatusCode, resp.Usage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionRelaysUpstreamError(t *testing.T) {
|
||||
tr, _ := newTransport(t, http.StatusUnauthorized, `{"error":{"message":"bad key"}}`, "application/json")
|
||||
req, _ := providers.ParseChatRequest([]byte(`{"model":"m","messages":[]}`))
|
||||
resp, err := tr.ChatCompletion(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusUnauthorized || string(resp.Body) != `{"error":{"message":"bad key"}}` {
|
||||
t.Errorf("status=%d body=%s", resp.StatusCode, resp.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionStreamSetsOptionsAndParsesEvents(t *testing.T) {
|
||||
sse := "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n" +
|
||||
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}\n\n" +
|
||||
"data: [DONE]\n\n"
|
||||
tr, got := newTransport(t, http.StatusOK, sse, "text/event-stream")
|
||||
req, _ := providers.ParseChatRequest([]byte(`{"model":"m","messages":[],"stream":true}`))
|
||||
|
||||
stream, err := tr.ChatCompletionStream(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer stream.Close()
|
||||
if got.headers.Get("Accept") != "text/event-stream" {
|
||||
t.Errorf("Accept = %q", got.headers.Get("Accept"))
|
||||
}
|
||||
if got.body["stream"] != true {
|
||||
t.Errorf("stream flag not set: %v", got.body)
|
||||
}
|
||||
if so, _ := got.body["stream_options"].(map[string]any); so["include_usage"] != true {
|
||||
t.Errorf("stream_options = %v", got.body["stream_options"])
|
||||
}
|
||||
n := 0
|
||||
for {
|
||||
_, err := stream.Recv()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
n++
|
||||
}
|
||||
if n != 2 {
|
||||
t.Errorf("events = %d, want 2 (the [DONE] marker is not an event)", n)
|
||||
}
|
||||
if stream.Usage().TotalTokens != 2 {
|
||||
t.Errorf("usage = %+v", stream.Usage())
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatCompletionStreamUpstreamErrorIsAPIError(t *testing.T) {
|
||||
tr, _ := newTransport(t, http.StatusTooManyRequests, `{"error":{"message":"slow down"}}`, "application/json")
|
||||
req, _ := providers.ParseChatRequest([]byte(`{"model":"m","messages":[],"stream":true}`))
|
||||
_, err := tr.ChatCompletionStream(context.Background(), req)
|
||||
var apiErr *providers.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %v, want *APIError", err)
|
||||
}
|
||||
if apiErr.StatusCode != http.StatusTooManyRequests || apiErr.Provider != "p" || string(apiErr.Body) != `{"error":{"message":"slow down"}}` {
|
||||
t.Errorf("apiErr = %+v", apiErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmbeddingsOverridesModel(t *testing.T) {
|
||||
tr, got := newTransport(t, http.StatusOK, `{"data":[],"usage":{"prompt_tokens":7,"total_tokens":7}}`, "application/json")
|
||||
req, _ := providers.ParseEmbeddingRequest([]byte(`{"model":"alias","input":"x"}`))
|
||||
req.Model = "real"
|
||||
resp, err := tr.Embeddings(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.path != "/v1/embeddings" || got.body["model"] != "real" || got.body["input"] != "x" {
|
||||
t.Errorf("path=%q body=%v", got.path, got.body)
|
||||
}
|
||||
if resp.Usage.TotalTokens != 7 {
|
||||
t.Errorf("usage = %+v", resp.Usage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModelsKeepsRawEntries(t *testing.T) {
|
||||
tr, got := newTransport(t, http.StatusOK,
|
||||
`{"data":[{"id":"a","owned_by":"x","extra":{"k":1}},{"id":"b"}]}`, "application/json")
|
||||
models, err := tr.ListModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.path != "/v1/models" {
|
||||
t.Errorf("path = %q", got.path)
|
||||
}
|
||||
if len(models) != 2 || models[0].ID != "a" || models[0].OwnedBy != "x" || models[1].ID != "b" {
|
||||
t.Errorf("models = %+v", models)
|
||||
}
|
||||
var raw map[string]any
|
||||
if err := json.Unmarshal(models[0].Raw, &raw); err != nil || raw["extra"] == nil {
|
||||
t.Errorf("raw entry lost: %s", models[0].Raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModelsUpstreamError(t *testing.T) {
|
||||
tr, _ := newTransport(t, http.StatusInternalServerError, `boom`, "text/plain")
|
||||
_, err := tr.ListModels(context.Background())
|
||||
var apiErr *providers.APIError
|
||||
if !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusInternalServerError {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBearerAuthNoopWhenEmpty(t *testing.T) {
|
||||
req, _ := http.NewRequest(http.MethodGet, "http://x", nil)
|
||||
if err := BearerAuth("")(context.Background(), req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if req.Header.Get("Authorization") != "" {
|
||||
t.Errorf("empty token must not set a header")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTransportUnreachableUpstream(t *testing.T) {
|
||||
srv := httptest.NewServer(http.NotFoundHandler())
|
||||
srv.Close() // closed on purpose
|
||||
tr := &Transport{Provider: "dead", BaseURL: srv.URL, HTTP: srv.Client()}
|
||||
req, _ := providers.ParseChatRequest([]byte(`{"model":"m"}`))
|
||||
if _, err := tr.ChatCompletion(context.Background(), req); err == nil {
|
||||
t.Fatal("expected a transport error")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user