Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
83 lines
3.2 KiB
Go
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"))
|
|
}
|
|
}
|