package channel import ( "context" "os" "testing" "time" "github.com/openteam/server/internal/pkg/crypto" "github.com/openteam/server/internal/store" "gorm.io/gorm" ) func newTestDB(t *testing.T) *gorm.DB { t.Helper() dsn := "/tmp/channel_test_" + time.Now().Format("150405.000000000") + ".db" db, err := store.Open("sqlite", dsn) if err != nil { t.Fatalf("open db: %v", err) } t.Cleanup(func() { sqlDB, _ := db.DB() sqlDB.Close() os.Remove(dsn) }) return db } func seedChannel(t *testing.T, db *gorm.DB, name string, weight, priority, maxConc int) uint64 { t.Helper() enc := crypto.NewEncryptor("test-master-key-1234567890") key, _ := enc.Encrypt("upstream-key") ch := store.Channel{ Name: name, Provider: store.ChannelProviderOpenAI, BaseURL: "http://localhost:9000", APIKeyEnc: key, Weight: weight, Priority: priority, MaxConcurrency: maxConc, HealthStatus: store.ChannelHealthHealthy, Enabled: true, } if err := db.Create(&ch).Error; err != nil { t.Fatalf("create channel: %v", err) } return ch.ID } func TestCandidatesFiltersUnhealthy(t *testing.T) { db := newTestDB(t) seedChannel(t, db, "a", 1, 0, 4) seedChannel(t, db, "b", 5, 1, 4) // 禁用渠道不入候选 bad := store.Channel{Name: "bad", Provider: store.ChannelProviderOpenAI, BaseURL: "http://x", APIKeyEnc: "x", HealthStatus: store.ChannelHealthCooldown, Enabled: true} db.Create(&bad) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) cands := s.Candidates("") if len(cands) != 2 { t.Fatalf("candidates = %d, want 2", len(cands)) } if cands[0].Channel.Name != "a" { t.Fatalf("first by priority should be a, got %s", cands[0].Channel.Name) } } func TestPickWeighted(t *testing.T) { db := newTestDB(t) seedChannel(t, db, "only", 1, 0, 4) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) cands := s.Candidates("") chosen := s.Pick(cands) if chosen == nil || chosen.Name != "only" { t.Fatalf("pick single: %+v", chosen) } // 空候选返回 nil if p := s.Pick(nil); p != nil { t.Fatal("expected nil for empty candidates") } } func TestTryAcquire(t *testing.T) { db := newTestDB(t) id := seedChannel(t, db, "c", 1, 0, 1) var ch store.Channel db.First(&ch, id) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) release, ok := s.TryAcquire(&ch) if !ok { t.Fatal("first acquire should succeed") } if _, ok := s.TryAcquire(&ch); ok { t.Fatal("second acquire should fail (capacity 1)") } release() if _, ok := s.TryAcquire(&ch); !ok { t.Fatal("acquire after release should succeed") } } func TestHealthCooldownAndRecovery(t *testing.T) { db := newTestDB(t) id := seedChannel(t, db, "h", 1, 0, 4) var ch store.Channel db.First(&ch, id) m := NewHealthMonitor(db, crypto.NewEncryptor("test-master-key-1234567890"), HealthConfig{ Interval: time.Minute, FailThreshold: 2, Timeout: time.Second, }) // 两次失败 → cooldown m.record(&ch, false) if ch.HealthStatus != store.ChannelHealthHealthy { t.Fatal("one failure should stay healthy") } m.record(&ch, false) if ch.HealthStatus != store.ChannelHealthCooldown { t.Fatalf("two failures should cooldown, got %s", ch.HealthStatus) } // 恢复 m.record(&ch, true) if ch.HealthStatus != store.ChannelHealthHealthy { t.Fatalf("success should recover, got %s", ch.HealthStatus) } } func TestHealthStartStops(t *testing.T) { db := newTestDB(t) seedChannel(t, db, "probe", 1, 0, 4) m := NewHealthMonitor(db, crypto.NewEncryptor("test-master-key-1234567890"), HealthConfig{ Interval: time.Hour, FailThreshold: 2, Timeout: time.Second, }) ctx, cancel := context.WithCancel(context.Background()) m.Start(ctx) cancel() // 不应 panic } // 模型被渠道绑定后只允许走绑定渠道;绑定渠道不可用时返回空候选(上层 503), // 不能退化到"不认识这个模型"的全局渠道。 func TestBoundModelDoesNotFallBackToGlobal(t *testing.T) { db := newTestDB(t) bound := seedChannel(t, db, "bound", 1, 0, 4) seedChannel(t, db, "other", 1, 1, 4) m := store.Model{Name: "gpt-m", Enabled: true} if err := db.Create(&m).Error; err != nil { t.Fatal(err) } if err := db.Create(&store.ChannelModelBinding{ChannelID: bound, ModelID: m.ID, UpstreamModel: "gpt-m"}).Error; err != nil { t.Fatal(err) } s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) cands := s.Candidates("gpt-m") if len(cands) != 1 || cands[0].Channel.Name != "bound" { t.Fatalf("绑定模型应只命中 bound,实际 %+v", cands) } if cands[0].UpstreamModel != "gpt-m" { t.Fatalf("映射名丢失:%s", cands[0].UpstreamModel) } db.Model(&store.Channel{}).Where("id = ?", bound).Update("enabled", false) if got := s.Candidates("gpt-m"); len(got) != 0 { t.Fatalf("绑定渠道停用后应为空候选,实际 %d", len(got)) } // 完全未绑定的模型仍然走全局 if got := s.Candidates("unmanaged"); len(got) != 1 || got[0].Channel.Name != "other" { t.Fatalf("未纳管模型应退回全局渠道,实际 %+v", got) } } // 运行期冷却:连续失败才踢,成功即清零,冷却到期自动恢复。 func TestRuntimeCooldownThresholdAndExpiry(t *testing.T) { db := newTestDB(t) id := seedChannel(t, db, "r", 1, 0, 4) var ch store.Channel db.First(&ch, id) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) s.cooldownFor = 60 * time.Millisecond // 默认阈值 3:前两次失败不该踢,第三次才进冷却 s.ReportFailure(&ch, 0) s.ReportFailure(&ch, 0) if s.InCooldown(id) { t.Fatal("未达阈值(3 次连续)不应进冷却") } s.ReportFailure(&ch, 0) if !s.InCooldown(id) { t.Fatal("连续达阈值应进冷却") } if got := s.Candidates(""); len(got) != 0 { t.Fatalf("冷却渠道不应进候选,实际 %d", len(got)) } if got := s.Candidates("whatever"); len(got) != 0 { t.Fatalf("全局候选也应过滤冷却渠道,实际 %d", len(got)) } time.Sleep(80 * time.Millisecond) if s.InCooldown(id) { t.Fatal("冷却到期应恢复可用") } if got := s.Candidates(""); len(got) != 1 { t.Fatalf("恢复后应重新进候选,实际 %d", len(got)) } s.ReportFailure(&ch, 0) s.ReportSuccess(&ch) s.ReportFailure(&ch, 0) if s.InCooldown(id) { t.Fatal("非连续失败不应触发冷却") } } // 管理后台可见性:RuntimeState 报累计失败与冷却剩余,ClearRuntime 立即清零。 func TestRuntimeStateAndClear(t *testing.T) { db := newTestDB(t) id := seedChannel(t, db, "s", 1, 0, 4) var ch store.Channel db.First(&ch, id) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) if fails, rem := s.RuntimeState(id); fails != 0 || rem != 0 { t.Fatalf("初始状态 = %d/%s,期望 0/0", fails, rem) } s.ReportFailure(&ch, 0) s.ReportFailure(&ch, 0) if fails, rem := s.RuntimeState(id); fails != 2 || rem != 0 { t.Fatalf("两次失败后 = %d/%s,期望 2/0", fails, rem) } s.ReportFailure(&ch, 0) fails, rem := s.RuntimeState(id) if fails != 0 || rem <= 0 || rem > s.cooldownFor { t.Fatalf("冷却中 = %d/%s,期望 0/(0,%s]", fails, rem, s.cooldownFor) } if got := len(s.Candidates("")); got != 0 { t.Fatalf("冷却渠道不应进候选,实际 %d", got) } s.ClearRuntime(id) if fails, rem := s.RuntimeState(id); fails != 0 || rem != 0 { t.Fatalf("清除后 = %d/%s,期望 0/0", fails, rem) } if s.InCooldown(id) { t.Fatal("清除后不应仍在冷却") } if got := len(s.Candidates("")); got != 1 { t.Fatalf("恢复后应重新进候选,实际 %d", got) } } // 上游 429 带的 Retry-After 可以拉长冷却,但要有上限,不能被钉死几小时。 func TestReportFailureCooldownOverrideAndClamp(t *testing.T) { db := newTestDB(t) id := seedChannel(t, db, "ra", 1, 0, 4) var ch store.Channel db.First(&ch, id) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) s.cooldownFor = 30 * time.Second for i := 0; i < 3; i++ { s.ReportFailure(&ch, 2*time.Minute) } if _, rem := s.RuntimeState(id); rem < 115*time.Second || rem > 2*time.Minute { t.Fatalf("应按 Retry-After 冷却 2min,实际剩余 %s", rem) } s.ClearRuntime(id) for i := 0; i < 3; i++ { s.ReportFailure(&ch, time.Hour) } if _, rem := s.RuntimeState(id); rem > maxCooldown { t.Fatalf("冷却应夹到上限 %s,实际 %s", maxCooldown, rem) } // 比默认值短的 Retry-After 不缩短冷却(否则 1s 就把坏渠道放回轮转) s.ClearRuntime(id) for i := 0; i < 3; i++ { s.ReportFailure(&ch, 5*time.Second) } if _, rem := s.RuntimeState(id); rem < 25*time.Second { t.Fatalf("短 Retry-After 不应缩短默认冷却,实际 %s", rem) } } // Pick 只在最高优先级层内加权随机;低优先级渠道即使权重再大也不该被首选打中。 func TestPickStaysInTopPriorityTier(t *testing.T) { db := newTestDB(t) seedChannel(t, db, "low-pri", 100000, 5, 4) // 低优先级 + 巨大权重 seedChannel(t, db, "top-pri", 1, 0, 4) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) cands := s.Candidates("") if len(cands) != 2 { t.Fatalf("candidates = %d, want 2", len(cands)) } for i := 0; i < 50; i++ { ch := s.Pick(cands) if ch.Name != "top-pri" { t.Fatalf("pick %d: 应始终选最高优先级渠道,实际 %s", i, ch.Name) } } } // 绑定级权重与渠道级权重相乘:默认绑定权重 1 时等价纯渠道权重,显式加权可单独放大。 func TestPickBindingWeightMultiplier(t *testing.T) { db := newTestDB(t) seedChannel(t, db, "plain", 1, 0, 4) seedChannel(t, db, "boosted", 1, 0, 4) var m store.Model m = store.Model{Name: "mm", Enabled: true} db.Create(&m) var chs []store.Channel db.Order("id").Find(&chs) db.Create(&store.ChannelModelBinding{ChannelID: chs[0].ID, ModelID: m.ID, UpstreamModel: "mm"}) // 默认 weight=1 db.Create(&store.ChannelModelBinding{ChannelID: chs[1].ID, ModelID: m.ID, UpstreamModel: "mm", Weight: 100}) // 显式放大 s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) cands := s.Candidates("mm") if len(cands) != 2 { t.Fatalf("candidates = %d, want 2", len(cands)) } boosted := 0 for i := 0; i < 200; i++ { if ch := s.Pick(cands); ch.Name == "boosted" { boosted++ } } if boosted < 150 || boosted > 200 { t.Fatalf("权重 100:1 下 boosted 命中 %d/200,应接近 200", boosted) } } // MaxConcurrency 改配置后在旧槽排空时热生效,无需重启。 func TestTryAcquireCapacityChange(t *testing.T) { db := newTestDB(t) id := seedChannel(t, db, "c", 1, 0, 1) var ch store.Channel db.First(&ch, id) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) release, ok := s.TryAcquire(&ch) if !ok { t.Fatal("容量 1 首次获取应成功") } // 容量 1 → 2:旧槽被占用,本次应沿用旧容量(拿不到) ch.MaxConcurrency = 2 if _, ok := s.TryAcquire(&ch); ok { t.Fatal("旧槽占用期间不应扩容") } release() // 排空后再取:新容量生效,连取 2 个都成功 r1, ok1 := s.TryAcquire(&ch) r2, ok2 := s.TryAcquire(&ch) if !ok1 || !ok2 { t.Fatal("排空后容量 2 应允许两个并发") } r1() r2() }