openteam/server/internal/channel/channel_test.go

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)
	}
}