Files
llmgw/internal/router/canary_test.go
T

83 lines
3.2 KiB
Go

package router
import (
"fmt"
"net/http"
"strings"
"testing"
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)
func canaryRouter(t *testing.T, percent int) (*Router, *fakeEngine, *fakeEngine) {
t.Helper()
stable := newFakeEngine(t, "qwen")
cand := newFakeEngine(t, "qwen-canary")
cfg := baseConfig(
config.Endpoint{ID: "s", URL: stable.srv.URL, Models: []string{"qwen"}},
config.Endpoint{ID: "c", URL: cand.srv.URL, Models: []string{"qwen-canary"}},
)
m := catModel("qwen", nil)
m.Canary = &catalog.Canary{Model: "qwen-canary", Percent: percent}
cfg.Models = []catalog.Model{m, catModel("qwen-canary", nil)}
return newTestRouter(t, cfg), stable, cand
}
func TestACanaryTakesItsShareAndIsComparedWithTheRest(t *testing.T) {
rt, stable, cand := canaryRouter(t, 30)
for i := range 400 {
rec := post(t, rt.Handler(), map[string]any{"model": "qwen", "messages": []any{}},
map[string]string{"X-Request-Id": fmt.Sprintf("r-%d", i)})
if rec.Code != http.StatusOK {
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
}
}
c, s := cand.calls.Load(), stable.calls.Load()
if c < 80 || c > 160 || c+s != 400 {
t.Errorf("canary %d, stable %d of 400 at 30%%", c, s)
}
if b, _ := cand.last.Load().([]byte); !strings.Contains(string(b), `"qwen-canary"`) {
t.Errorf("the candidate was asked for another name: %s", b)
}
st := rt.CanaryStats()["qwen"]
if st.Canary != "qwen-canary" || st.Candidate.Requests != c || st.Stable.Requests != s {
t.Errorf("stats %+v", st)
}
}
func TestASessionStaysOnOneSideAndAHeaderPinsIt(t *testing.T) {
rt, _, cand := canaryRouter(t, 50)
for range 20 {
post(t, rt.Handler(), map[string]any{"model": "qwen", "user": "alice", "messages": []any{}}, nil)
}
if n := cand.calls.Load(); n != 0 && n != 20 {
t.Errorf("one session went to the canary %d times of 20", n)
}
before := cand.calls.Load()
rec := post(t, rt.Handler(), map[string]any{"model": "qwen", "user": "alice", "messages": []any{}},
map[string]string{"X-Llmgw-Variant": "canary"})
if cand.calls.Load() != before+1 || rec.Header().Get("x-llmgw-variant") != "canary" {
t.Errorf("the pin was not followed: %q", rec.Header().Get("x-llmgw-variant"))
}
}
func TestACanaryCannotBeNamedAndFallsBackToStable(t *testing.T) {
rt, _, _ := canaryRouter(t, 100)
if rec := post(t, rt.Handler(), map[string]any{"model": "qwen-canary", "messages": []any{}}, nil); rec.Code != http.StatusNotFound {
t.Errorf("a canary named directly: %d", rec.Code)
}
// With nothing ready on the canary side, the stable side serves.
stable := newFakeEngine(t, "qwen")
cfg := baseConfig(config.Endpoint{ID: "s", URL: stable.srv.URL, Models: []string{"qwen"}})
m := catModel("qwen", nil)
m.Canary = &catalog.Canary{Model: "qwen-canary", Percent: 100}
cfg.Models = []catalog.Model{m, catModel("qwen-canary", nil)}
rt = newTestRouter(t, cfg)
rec := post(t, rt.Handler(), map[string]any{"model": "qwen", "messages": []any{}}, nil)
if rec.Code != http.StatusOK || stable.calls.Load() != 1 || rec.Header().Get("x-llmgw-variant") != "stable" {
t.Errorf("status %d, stable calls %d, variant %q", rec.Code, stable.calls.Load(), rec.Header().Get("x-llmgw-variant"))
}
}