internal/pool/modelview_test.go

package pool

import (
	"testing"
	"time"

	"github.com/linguo2625469/workbuddy2api-panel/internal/auth"
)

// ---------------------------------------------------------------------------
// ModelLockView：把「哪些模型不能用、锁了几个号、还要锁多久」聚合成一张全清单
// （ModelBlocked 是单模型按请求回答，本视图是全清单，口径必须与选号一致）
// ---------------------------------------------------------------------------

// lockViewPool 造一个含两个 cn 账号的池，便于手搓模型级冷却。
func lockViewPool(t *testing.T) *Pool {
	t.Helper()
	withNoPickGap(t)
	p := New("")
	p.Add(&auth.Auth{UID: "u1"})
	p.Add(&auth.Auth{UID: "u2"})
	return p
}

// setModelCooldown 直接写模型级冷却（与 modelblocked_test.go 同法）。
func setModelCooldown(p *Pool, uid, model string, until time.Time, reason string, auditOnly bool) {
	p.mu.Lock()
	e := p.byUID[uid]
	if e.modelCooldowns == nil {
		e.modelCooldowns = map[string]modelCooldown{}
	}
	e.modelCooldowns[model] = modelCooldown{Until: until, Reason: reason, AuditOnly: auditOnly}
	p.mu.Unlock()
}

// TestModelLockViewEmpty 没有未过期模型级冷却 → 返回 nil（面板按空态渲染）。
func TestModelLockViewEmpty(t *testing.T) {
	p := lockViewPool(t)
	if rows := p.ModelLockView(); rows != nil {
		t.Fatalf("无冷却时应返回 nil，得到 %+v", rows)
	}
}

// TestModelLockViewPartial 部分号被锁 → partial，且带出可选/总数、最早解锁与原因。
func TestModelLockViewPartial(t *testing.T) {
	p := lockViewPool(t)
	setModelCooldown(p, "u1", "glm-5.3", time.Now().Add(time.Hour), "6004 model rate limit", false)

	rows := p.ModelLockView()
	if len(rows) != 1 {
		t.Fatalf("want 1 row, got %+v", rows)
	}
	r := rows[0]
	if r.Model != "glm-5.3" || r.Realm != "cn" {
		t.Errorf("model/realm = %q/%q want glm-5.3/cn", r.Model, r.Realm)
	}
	if r.Total != 2 || r.Locked != 1 || r.Servable != 1 {
		t.Errorf("total/locked/servable = %d/%d/%d want 2/1/1", r.Total, r.Locked, r.Servable)
	}
	if r.State != "partial" {
		t.Errorf("state = %q want partial", r.State)
	}
	if r.Reason != "6004 model rate limit" {
		t.Errorf("reason = %q want 6004 model rate limit", r.Reason)
	}
	if d := time.Until(r.UnlockAt); d < 50*time.Minute || d > 70*time.Minute {
		t.Errorf("unlock_at 应在约 1h 后，得到 %v", d)
	}
	if !r.FullyUnlockAt.Equal(r.UnlockAt) {
		t.Errorf("只有一个号被锁时 fully_unlock_at 应等于 unlock_at")
	}
}

// TestModelLockViewLocked 全部参与选号的号都被该模型挡住 → locked（与 ModelBlocked 同义），
// 最早解锁取最小 Until、全池解锁取最大 Until，原因取最早解锁账号的。
func TestModelLockViewLocked(t *testing.T) {
	p := lockViewPool(t)
	base := time.Now()
	setModelCooldown(p, "u1", "glm-5.3", base.Add(2*time.Hour), "reason-late", false)
	setModelCooldown(p, "u2", "glm-5.3", base.Add(time.Hour), "reason-early", false)

	rows := p.ModelLockView()
	if len(rows) != 1 {
		t.Fatalf("want 1 row, got %+v", rows)
	}
	r := rows[0]
	if r.State != "locked" {
		t.Errorf("state = %q want locked", r.State)
	}
	if r.Servable != 0 || r.Locked != 2 || r.Total != 2 {
		t.Errorf("servable/locked/total = %d/%d/%d want 0/2/2", r.Servable, r.Locked, r.Total)
	}
	if r.Reason != "reason-early" {
		t.Errorf("reason 应取最早解锁账号的，得到 %q", r.Reason)
	}
	if d := time.Until(r.UnlockAt); d < 50*time.Minute || d > 70*time.Minute {
		t.Errorf("unlock_at 应约 1h 后，得到 %v", d)
	}
	if d := time.Until(r.FullyUnlockAt); d < 110*time.Minute || d > 130*time.Minute {
		t.Errorf("fully_unlock_at 应约 2h 后，得到 %v", d)
	}
}

// TestModelLockViewStarved 此刻没号能服务、但不是模型冷却造成的（账号级冷却）→ starved，
// 与 locked 区分开：这种情况换模型没用，等一下就好。
func TestModelLockViewStarved(t *testing.T) {
	p := lockViewPool(t)
	setModelCooldown(p, "u1", "glm-5.3", time.Now().Add(time.Hour), "6004 model rate limit", false)
	p.Cooldown("u2", CoolHard, time.Hour, "14018 credits exhausted")

	rows := p.ModelLockView()
	if len(rows) != 1 {
		t.Fatalf("want 1 row, got %+v", rows)
	}
	r := rows[0]
	if r.Servable != 0 || r.Locked != 1 || r.Total != 2 {
		t.Errorf("servable/locked/total = %d/%d/%d want 0/1/2", r.Servable, r.Locked, r.Total)
	}
	if r.State != "starved" {
		t.Errorf("state = %q want starved（模型没全锁，只是此刻没号）", r.State)
	}
}

// TestModelLockViewSkipsAuditOnly AuditOnly 台账只审计不拦路由 → 既不产生行，
// 也不计入 locked（口径与 modelCooled 一致，否则会把其实能调的模型报成不可用）。
func TestModelLockViewSkipsAuditOnly(t *testing.T) {
	p := lockViewPool(t)
	setModelCooldown(p, "u1", "glm-5.3", time.Now().Add(time.Hour), "audit", true)
	if rows := p.ModelLockView(); rows != nil {
		t.Fatalf("AuditOnly 条目不应产生行：%+v", rows)
	}
}

// TestModelLockViewSkipsExpired 已过期的冷却不算锁。
func TestModelLockViewSkipsExpired(t *testing.T) {
	p := lockViewPool(t)
	setModelCooldown(p, "u1", "glm-5.3", time.Now().Add(-time.Minute), "6004", false)
	if rows := p.ModelLockView(); rows != nil {
		t.Fatalf("已过期冷却不应产生行：%+v", rows)
	}
}

// TestModelLockViewSkipsDisabledAndPaused 禁用号/暂停号的不可用与模型无关：
// 既不计入分母 Total，也不该被算成「还能服务的证明」。
func TestModelLockViewSkipsDisabledAndPaused(t *testing.T) {
	p := lockViewPool(t)
	setModelCooldown(p, "u1", "glm-5.3", time.Now().Add(time.Hour), "6004", false)
	p.mu.Lock()
	p.byUID["u2"].disabled = true
	p.mu.Unlock()

	rows := p.ModelLockView()
	if len(rows) != 1 {
		t.Fatalf("want 1 row, got %+v", rows)
	}
	if rows[0].Total != 1 {
		t.Errorf("total = %d want 1（禁用号不计入分母）", rows[0].Total)
	}
	if rows[0].State != "locked" {
		t.Errorf("state = %q want locked（唯一参与选号的号被锁）", rows[0].State)
	}

	p.mu.Lock()
	p.byUID["u2"].disabled = false
	p.byUID["u2"].paused = true
	p.mu.Unlock()

	rows = p.ModelLockView()
	if len(rows) != 1 || rows[0].Total != 1 {
		t.Errorf("暂停号同样不计入分母：%+v", rows)
	}
}

// TestModelLockViewRealmIsolation 同名模型在 cn/global 各自独立计数与原因。
func TestModelLockViewRealmIsolation(t *testing.T) {
	p := realmPool(t)
	setModelCooldown(p, "cn1", "glm-5.3", time.Now().Add(time.Hour), "6004 cn", false)
	setModelCooldown(p, "g1", "glm-5.3", time.Now().Add(time.Hour), "6004 global", false)

	rows := p.ModelLockView()
	if len(rows) != 2 {
		t.Fatalf("同名模型在两个域应各出一行，得到 %+v", rows)
	}
	byRealm := map[string]ModelLockRow{}
	for _, r := range rows {
		byRealm[r.Realm] = r
	}
	cn, ok := byRealm["cn"]
	if !ok || cn.Total != 2 || cn.Locked != 1 || cn.Reason != "6004 cn" {
		t.Errorf("cn 行 = %+v want total 2 locked 1 reason \"6004 cn\"", cn)
	}
	g, ok := byRealm["global"]
	if !ok || g.Total != 2 || g.Locked != 1 || g.Reason != "6004 global" {
		t.Errorf("global 行 = %+v want total 2 locked 1 reason \"6004 global\"", g)
	}
}

// TestModelLockViewSortsLockedFirst 整池不可用排在部分限流前面，方便一眼看到最该处理的。
func TestModelLockViewSortsLockedFirst(t *testing.T) {
	p := lockViewPool(t)
	base := time.Now()
	setModelCooldown(p, "u1", "glm-5.3", base.Add(time.Hour), "6004", false)
	setModelCooldown(p, "u2", "glm-5.3", base.Add(time.Hour), "6004", false)
	setModelCooldown(p, "u1", "kimi-k2", base.Add(time.Hour), "6004", false)

	rows := p.ModelLockView()
	if len(rows) != 2 {
		t.Fatalf("want 2 rows, got %+v", rows)
	}
	if rows[0].State != "locked" || rows[1].State != "partial" {
		t.Errorf("整池不可用应排在部分限流前：%+v", rows)
	}
	if rows[0].Model != "glm-5.3" {
		t.Errorf("首行应为 glm-5.3，得到 %q", rows[0].Model)
	}
}

// TestModelLockViewAfterClear 走真实写入/清除通道（而非只手搓 map）：
// BlockModelBackoff 后应出 partial 行，BlockModelClear 后应回到 nil。
func TestModelLockViewAfterClear(t *testing.T) {
	p := lockViewPool(t)
	p.BlockModelBackoff("u1", "glm-5.3", "11102 model [glm-5.3] service info not found")
	rows := p.ModelLockView()
	if len(rows) != 1 || rows[0].State != "partial" {
		t.Fatalf("写入负缓存后应有一行 partial：%+v", rows)
	}
	p.BlockModelClear("u1", "glm-5.3")
	if rows := p.ModelLockView(); rows != nil {
		t.Fatalf("清除后应回到 nil：%+v", rows)
	}
}