1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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)
}
}
|
