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.
237 lines
8.1 KiB
Go
237 lines
8.1 KiB
Go
package store_test
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
|
)
|
|
|
|
// TestStoreQueries exercises the hand-written SQL against a real Postgres with
|
|
// migrations applied. Set GATEWAY_TEST_DB_URL to run.
|
|
func TestStoreQueries(t *testing.T) {
|
|
url := os.Getenv("GATEWAY_TEST_DB_URL")
|
|
if url == "" {
|
|
t.Skip("set GATEWAY_TEST_DB_URL to run the store DB integration test")
|
|
}
|
|
ctx := context.Background()
|
|
pool, err := pgxpool.New(ctx, url)
|
|
if err != nil {
|
|
t.Fatalf("connect: %v", err)
|
|
}
|
|
defer pool.Close()
|
|
for _, table := range []string{"request_logs", "virtual_keys", "providers", "users"} {
|
|
if _, err := pool.Exec(ctx, "DELETE FROM "+table); err != nil {
|
|
t.Fatalf("cleanup %s: %v", table, err)
|
|
}
|
|
}
|
|
st := store.New(pool)
|
|
|
|
// Empty window: every aggregate must scan as zero, not NULL.
|
|
sum, err := st.GetLogsSummary(ctx, time.Now().Add(-time.Hour), nil)
|
|
if err != nil {
|
|
t.Fatalf("GetLogsSummary on empty table: %v", err)
|
|
}
|
|
if sum.Requests != 0 || sum.Errors != 0 || sum.TotalCost != 0 {
|
|
t.Fatalf("empty summary = %+v", sum)
|
|
}
|
|
|
|
user, err := st.UpsertUser(ctx, "sub-1", "a@example.com", "A")
|
|
if err != nil {
|
|
t.Fatalf("UpsertUser: %v", err)
|
|
}
|
|
|
|
// Keys: rules and created_at come back from the list query.
|
|
keyID, err := st.CreateVirtualKey(ctx, store.CreateVirtualKeyParams{
|
|
Name: "k", KeyHash: []byte("hash-1"), KeyPrefix: "sk-00000001", UserID: &user.ID,
|
|
Rules: []string{"local/*", "openai/gpt-4o"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateVirtualKey: %v", err)
|
|
}
|
|
keys, err := st.ListVirtualKeys(ctx, &user.ID)
|
|
if err != nil {
|
|
t.Fatalf("ListVirtualKeys: %v", err)
|
|
}
|
|
if len(keys) != 1 || len(keys[0].Rules) != 2 || keys[0].Rules[0] != "local/*" || keys[0].CreatedAt.IsZero() {
|
|
t.Fatalf("keys = %+v", keys)
|
|
}
|
|
|
|
// Providers: update, counts, and the catalog join.
|
|
provID, err := st.CreateProvider(ctx, store.CreateProviderParams{Name: "local", Type: "llamacpp", Enabled: true})
|
|
if err != nil {
|
|
t.Fatalf("CreateProvider: %v", err)
|
|
}
|
|
newURL := "http://10.0.0.5:8080/v1"
|
|
off := false
|
|
if err := st.UpdateProvider(ctx, provID, store.UpdateProviderParams{BaseURL: &newURL, Enabled: &off, Config: []byte(`{"a":"b"}`)}); err != nil {
|
|
t.Fatalf("UpdateProvider: %v", err)
|
|
}
|
|
// Other DB tests may run against the same database, so look the row up by id.
|
|
findProv := func() store.ProviderRow {
|
|
provs, err := st.ListProviders(ctx, false)
|
|
if err != nil {
|
|
t.Fatalf("ListProviders: %v", err)
|
|
}
|
|
for _, p := range provs {
|
|
if p.ID == provID {
|
|
return p
|
|
}
|
|
}
|
|
t.Fatalf("provider %s missing from %+v", provID, provs)
|
|
return store.ProviderRow{}
|
|
}
|
|
if p := findProv(); p.BaseURL != newURL || p.Enabled || string(p.Config) != `{"a": "b"}` {
|
|
t.Fatalf("provider after update = %+v", p)
|
|
}
|
|
empty := ""
|
|
if err := st.UpdateProvider(ctx, provID, store.UpdateProviderParams{BaseURL: &empty}); err != nil {
|
|
t.Fatalf("UpdateProvider clear url: %v", err)
|
|
}
|
|
if p := findProv(); p.BaseURL != "" {
|
|
t.Fatalf("base_url not cleared: %q", p.BaseURL)
|
|
}
|
|
|
|
ctxLen := 4096
|
|
if err := st.UpsertDiscovered(ctx, provID, time.Now(), []store.DiscoveredModel{
|
|
{Slug: "m1", ContextLength: &ctxLen}, {Slug: "m2"},
|
|
}); err != nil {
|
|
t.Fatalf("UpsertDiscovered: %v", err)
|
|
}
|
|
if err := st.UpsertDiscovered(ctx, provID, time.Now(), []store.DiscoveredModel{{Slug: "m1"}}); err != nil {
|
|
t.Fatalf("UpsertDiscovered 2: %v", err)
|
|
}
|
|
counts, err := st.ProviderModelCounts(ctx)
|
|
if err != nil {
|
|
t.Fatalf("ProviderModelCounts: %v", err)
|
|
}
|
|
if c := counts[provID]; c.Models != 2 || c.Available != 1 || c.LastSyncedAt == nil {
|
|
t.Fatalf("counts = %+v", c)
|
|
}
|
|
|
|
// Overrides: set, then clear with an explicit null patch.
|
|
rows, err := st.ListModelsAdmin(ctx)
|
|
if err != nil {
|
|
t.Fatalf("ListModelsAdmin: %v", err)
|
|
}
|
|
var m1, m2 uuid.UUID
|
|
for _, r := range rows {
|
|
if r.ProviderName == "local" && r.Slug == "m1" {
|
|
m1 = r.ID
|
|
}
|
|
if r.ProviderName == "local" && r.Slug == "m2" {
|
|
m2 = r.ID
|
|
}
|
|
}
|
|
if m1 == uuid.Nil || m2 == uuid.Nil {
|
|
t.Fatalf("local/m1 or local/m2 missing from %+v", rows)
|
|
}
|
|
price := 2.5
|
|
rpm := 30
|
|
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{
|
|
InputPrice: store.Patch[float64]{Set: true, Value: &price},
|
|
RPMLimit: store.Patch[int]{Set: true, Value: &rpm},
|
|
}); err != nil {
|
|
t.Fatalf("UpdateModelOverride set: %v", err)
|
|
}
|
|
rows, _ = st.ListModelsAdmin(ctx)
|
|
for _, r := range rows {
|
|
if r.ID == m1 && (r.PriceInputPerMTok == nil || *r.PriceInputPerMTok != 2.5 || r.RPMLimit == nil || *r.RPMLimit != 30) {
|
|
t.Fatalf("override not applied: %+v", r)
|
|
}
|
|
}
|
|
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{InputPrice: store.Patch[float64]{Set: true}}); err != nil {
|
|
t.Fatalf("UpdateModelOverride clear: %v", err)
|
|
}
|
|
rows, _ = st.ListModelsAdmin(ctx)
|
|
for _, r := range rows {
|
|
if r.ID == m1 {
|
|
if r.PriceInputPerMTok != nil {
|
|
t.Fatalf("price not cleared: %+v", r)
|
|
}
|
|
if r.RPMLimit == nil || *r.RPMLimit != 30 {
|
|
t.Fatalf("unset patch must leave rpm_limit alone: %+v", r)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Backup model: set, read back through the admin list and the catalog,
|
|
// reject an unknown target, then clear.
|
|
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{FallbackModelID: store.Patch[uuid.UUID]{Set: true, Value: &m2}}); err != nil {
|
|
t.Fatalf("UpdateModelOverride fallback: %v", err)
|
|
}
|
|
rows, _ = st.ListModelsAdmin(ctx)
|
|
for _, r := range rows {
|
|
if r.ID == m1 && (r.FallbackModelID == nil || *r.FallbackModelID != m2 || r.FallbackAddress == nil || *r.FallbackAddress != "local/m2") {
|
|
t.Fatalf("fallback not applied: %+v", r)
|
|
}
|
|
}
|
|
on := true
|
|
if err := st.UpdateProvider(ctx, provID, store.UpdateProviderParams{Enabled: &on}); err != nil {
|
|
t.Fatalf("UpdateProvider enable: %v", err)
|
|
}
|
|
cat, err := st.ListCatalog(ctx)
|
|
if err != nil {
|
|
t.Fatalf("ListCatalog: %v", err)
|
|
}
|
|
var sawM1 bool
|
|
for _, r := range cat {
|
|
if r.ID == m1 {
|
|
sawM1 = true
|
|
if r.FallbackModelID == nil || *r.FallbackModelID != m2 {
|
|
t.Fatalf("catalog row lacks fallback: %+v", r)
|
|
}
|
|
}
|
|
}
|
|
if !sawM1 {
|
|
t.Fatalf("local/m1 missing from catalog %+v", cat)
|
|
}
|
|
unknown := uuid.New()
|
|
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{FallbackModelID: store.Patch[uuid.UUID]{Set: true, Value: &unknown}}); err == nil {
|
|
t.Fatal("UpdateModelOverride must reject an unknown backup model")
|
|
}
|
|
if err := st.UpdateModelOverride(ctx, m1, store.ModelOverride{FallbackModelID: store.Patch[uuid.UUID]{Set: true}}); err != nil {
|
|
t.Fatalf("UpdateModelOverride clear fallback: %v", err)
|
|
}
|
|
rows, _ = st.ListModelsAdmin(ctx)
|
|
for _, r := range rows {
|
|
if r.ID == m1 && (r.FallbackModelID != nil || r.FallbackAddress != nil) {
|
|
t.Fatalf("fallback not cleared: %+v", r)
|
|
}
|
|
}
|
|
|
|
// Logs: insert, then list and summarize.
|
|
logID, err := st.InsertRequestLog(ctx, store.RequestLog{
|
|
VirtualKeyID: &keyID, UserID: &user.ID, ProviderID: &provID, ProviderModelID: &m2,
|
|
KeyPrefix: "sk-00000001", ModelName: "local/m1", ServedModel: "local/m2", ProviderName: "local",
|
|
StatusCode: 502, Error: "upstream down", LatencyMS: 12, ClientIP: "127.0.0.1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("InsertRequestLog: %v", err)
|
|
}
|
|
if err := st.InsertLogContent(ctx, logID, []byte(`[{"role":"user","content":"hi"}]`), ""); err != nil {
|
|
t.Fatalf("InsertLogContent: %v", err)
|
|
}
|
|
list, err := st.ListRequestLogs(ctx, store.ListLogsParams{UserID: &user.ID, Query: "local/m1", StatusGte: 400})
|
|
if err != nil || len(list) != 1 || list[0].Error != "upstream down" || list[0].ServedModel != "local/m2" {
|
|
t.Fatalf("ListRequestLogs: %v %+v", err, list)
|
|
}
|
|
// The search also matches the address that served the request.
|
|
if list, err := st.ListRequestLogs(ctx, store.ListLogsParams{UserID: &user.ID, Query: "local/m2"}); err != nil || len(list) != 1 {
|
|
t.Fatalf("ListRequestLogs by served model: %v %+v", err, list)
|
|
}
|
|
sum, err = st.GetLogsSummary(ctx, time.Now().Add(-time.Hour), &user.ID)
|
|
if err != nil || sum.Requests != 1 || sum.Errors != 1 {
|
|
t.Fatalf("summary = %+v err=%v", sum, err)
|
|
}
|
|
series, err := st.RequestVolumeSeries(ctx, time.Now().Add(-time.Hour), &user.ID, 6)
|
|
if err != nil || len(series) != 6 || series[5] != 1 {
|
|
t.Fatalf("series = %v err=%v", series, err)
|
|
}
|
|
}
|