summaryrefslogtreecommitdiff
path: root/internal/routing/router_test.go
diff options
context:
space:
mode:
Diffstat (limited to 'internal/routing/router_test.go')
-rw-r--r--internal/routing/router_test.go60
1 files changed, 60 insertions, 0 deletions
diff --git a/internal/routing/router_test.go b/internal/routing/router_test.go
new file mode 100644
index 0000000..62dc656
--- /dev/null
+++ b/internal/routing/router_test.go
@@ -0,0 +1,60 @@
+package routing
+
+import (
+ "testing"
+
+ "aigw/internal/catalog"
+ "aigw/internal/config"
+ "aigw/internal/domain"
+)
+
+func TestPlanHonorsPriorityAndProtocol(t *testing.T) {
+ cfg := config.Config{
+ Providers: []config.ProviderConfig{
+ {ID: "openai-primary", Protocol: domain.ProtocolOpenAI, BaseURL: "https://one.test", APIKey: "one"},
+ {ID: "openai-fallback", Protocol: domain.ProtocolOpenAI, BaseURL: "https://two.test", APIKey: "two"},
+ {ID: "anthropic", Protocol: domain.ProtocolAnthropic, BaseURL: "https://three.test", APIKey: "three"},
+ },
+ Models: []config.ModelConfig{{
+ ID: "public/model",
+ Routes: []config.RouteConfig{
+ {Provider: "openai-fallback", UpstreamModel: "fallback", Priority: 10, Weight: 1},
+ {Provider: "anthropic", UpstreamModel: "claude", Priority: 0, Weight: 1},
+ {Provider: "openai-primary", UpstreamModel: "primary", Priority: 0, Weight: 1},
+ },
+ }},
+ }
+ router := New(catalog.New(cfg))
+ plan, err := router.Plan("public/model", domain.ProtocolOpenAI)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if len(plan) != 2 || plan[0].Provider.ID != "openai-primary" || plan[1].Provider.ID != "openai-fallback" {
+ t.Fatalf("unexpected plan: %+v", plan)
+ }
+}
+
+func TestPlanUsesWeightsForPrimarySelection(t *testing.T) {
+ cfg := config.Config{
+ Providers: []config.ProviderConfig{
+ {ID: "one", Protocol: domain.ProtocolOpenAI, BaseURL: "https://one.test", APIKey: "one"},
+ {ID: "two", Protocol: domain.ProtocolOpenAI, BaseURL: "https://two.test", APIKey: "two"},
+ },
+ Models: []config.ModelConfig{{ID: "public/model", Routes: []config.RouteConfig{
+ {Provider: "one", UpstreamModel: "one", Weight: 3},
+ {Provider: "two", UpstreamModel: "two", Weight: 1},
+ }}},
+ }
+ router := New(catalog.New(cfg))
+ counts := map[string]int{}
+ for range 8 {
+ plan, err := router.Plan("public/model", domain.ProtocolOpenAI)
+ if err != nil {
+ t.Fatal(err)
+ }
+ counts[plan[0].Provider.ID]++
+ }
+ if counts["one"] != 6 || counts["two"] != 2 {
+ t.Fatalf("unexpected weighted distribution: %+v", counts)
+ }
+}