ModelRouter/internal/core/channels_pool_test.go

package core

import (
	"strings"
	"testing"

	"llmsproxy/internal/lua"

	"fmt"
	"net/http"
	"net/http/httptest"
)

// channels_pool_test：对齐参考 account.* / model.* / provider.* 契约的行为测试。
// 账号池增删改、上游源随 enabled/模型开关同步、供应商整体开关、备份导出。

func poolCore(t *testing.T) (*Core, *lua.Plugins) {
	t.Helper()
	c := kernelTestCore(t)
	ps := c.Plugins()
	src, err := readBundledChannels(t)
	if err != nil {
		t.Fatal(err)
	}
	if err := ps.LoadSource("channels", src); err != nil {
		t.Fatal(err)
	}
	_ = ps.SetPluginKernelPermission("channels", []string{"add_source", "list_sources", "remove_source", "status", "probe_all", "reset_health"}, false)
	return c, ps
}

// listSourcesViaKernel 直查内核 list_sources（actionListSources 返回 []map）。
func listSourcesViaKernel(t *testing.T, c *Core) []map[string]interface{} {
	t.Helper()
	out, err := c.Invoke(nil, "list_sources", map[string]interface{}{})
	if err != nil {
		t.Fatalf("list_sources: %v", err)
	}
	if ms, ok := out.([]map[string]interface{}); ok {
		return ms
	}
	raw, _ := out.([]interface{})
	var srcs []map[string]interface{}
	for _, s := range raw {
		if m, ok := s.(map[string]interface{}); ok {
			srcs = append(srcs, m)
		}
	}
	return srcs
}

func invoke(t *testing.T, ps *lua.Plugins, fn, body string) map[string]interface{} {
	t.Helper()
	out, err := ps.InvokeHandler("channels", fn, &lua.HandlerRequest{Method: "POST", Body: body})
	if err != nil {
		t.Fatalf("invoke %s: %v", fn, err)
	}
	return out.JSON.(map[string]interface{})
}

// asArr 规整桥接后的数组。Lua 空表过桥是 map[string]interface{}（无法与空数组
// 区分，见 vm.go luaTableToGoValue）；内核 list_sources 的 models 过桥是 []string。
func asArr(v interface{}) []interface{} {
	switch a := v.(type) {
	case []interface{}:
		return a
	case []string:
		out := make([]interface{}, len(a))
		for i, s := range a {
			out[i] = s
		}
		return out
	}
	return nil
}

func srcHasAccount(srcs []map[string]interface{}, needle string) (map[string]interface{}, bool) {
	for _, sm := range srcs {
		if name, _ := sm["name"].(string); strings.Contains(name, needle) {
			return sm, true
		}
	}
	return nil, false
}

// TestChannelsAccountPoolBehavior：账号池 CRUD + 上游源同步 + 模型/供应商开关 + 备份。
func TestChannelsAccountPoolBehavior(t *testing.T) {
	c, ps := poolCore(t)

	t.Run("账号增删改", func(t *testing.T) {
		m := invoke(t, ps, "account_add", `{"platform":"minimax","nickname":"我的号"}`)
		if m["ok"] != true {
			t.Fatalf("account_add: %v", m)
		}
		aid, _ := m["account_id"].(string)
		if aid == "" {
			t.Fatalf("account_add 应返回 account_id: %v", m)
		}
		m = invoke(t, ps, "account_list", `{"platform":"minimax"}`)
		accts := asArr(m["accounts"])
		if len(accts) != 1 {
			t.Fatalf("应有1个账号: %d", len(accts))
		}
		a := accts[0].(map[string]interface{})
		if a["nickname"] != "我的号" || a["enabled"] != true {
			t.Fatalf("账号字段: %v", a)
		}
		invoke(t, ps, "account_update", `{"platform":"minimax","account_id":"`+aid+`","patch":{"nickname":"改名"}}`)
		m = invoke(t, ps, "account_list", `{"platform":"minimax"}`)
		if asArr(m["accounts"])[0].(map[string]interface{})["nickname"] != "改名" {
			t.Fatalf("改名失败")
		}
		invoke(t, ps, "account_delete", `{"platform":"minimax","account_id":"`+aid+`"}`)
		m = invoke(t, ps, "account_list", `{"platform":"minimax"}`)
		if n := len(asArr(m["accounts"])); n != 0 {
			t.Fatalf("删除后应为0, 实得 %d", n)
		}
	})

	t.Run("模型开关同步上游源", func(t *testing.T) {
		m := invoke(t, ps, "account_add", `{"platform":"minimax","credential":"MM-KEY"}`)
		aid, _ := m["account_id"].(string)
		sm, ok := srcHasAccount(listSourcesViaKernel(t, c), aid)
		if !ok {
			t.Fatal("带凭据账号应已同步源")
		}
		if n := len(asArr(sm["models"])); n != 3 {
			t.Fatalf("源应有3模型, 实得 %d", n)
		}
		invoke(t, ps, "model_set_disabled", `{"platform":"minimax","model_id":"MiniMax-M3","disabled":true}`)
		sm, _ = srcHasAccount(listSourcesViaKernel(t, c), aid)
		if n := len(asArr(sm["models"])); n != 2 {
			t.Fatalf("关1模型后源应2模型, 实得 %d", n)
		}
		invoke(t, ps, "model_set_all_disabled", `{"platform":"minimax","disabled":true}`)
		m = invoke(t, ps, "provider_status", `{}`)
		for _, s := range asArr(m["statuses"]) {
			x := s.(map[string]interface{})
			if x["id"] == "minimax" && x["closed"] != true {
				t.Fatalf("全关模型后应 closed=true: %v", x)
			}
		}
		invoke(t, ps, "model_set_all_disabled", `{"platform":"minimax","disabled":false}`)
	})

	t.Run("enabled开关同步上游源", func(t *testing.T) {
		m := invoke(t, ps, "account_list", `{"platform":"minimax"}`)
		accts := asArr(m["accounts"])
		if len(accts) == 0 {
			t.Fatal("上个用例应留有账号")
		}
		aid := accts[0].(map[string]interface{})["id"].(string)
		before := len(listSourcesViaKernel(t, c))
		invoke(t, ps, "account_update", `{"platform":"minimax","account_id":"`+aid+`","patch":{"enabled":false}}`)
		if n := len(listSourcesViaKernel(t, c)); n != before-1 {
			t.Fatalf("停用账号应删源: 前 %d 后 %d", before, n)
		}
		invoke(t, ps, "account_update", `{"platform":"minimax","account_id":"`+aid+`","patch":{"enabled":true}}`)
		if n := len(listSourcesViaKernel(t, c)); n != before {
			t.Fatalf("启用账号应重建源: 前 %d 后 %d", before, n)
		}
	})

	t.Run("供应商整体开关", func(t *testing.T) {
		invoke(t, ps, "provider_set_enabled", `{"platform":"minimax","enabled":false}`)
		if _, ok := srcHasAccount(listSourcesViaKernel(t, c), "minimax"); ok {
			t.Fatal("关闭渠道后不应有 minimax 源")
		}
		m := invoke(t, ps, "provider_status", `{}`)
		for _, s := range asArr(m["statuses"]) {
			x := s.(map[string]interface{})
			if x["id"] == "minimax" && x["closed"] != true {
				t.Fatalf("关闭后应 closed: %v", x)
			}
		}
		invoke(t, ps, "provider_set_enabled", `{"platform":"minimax","enabled":true}`)
		if _, ok := srcHasAccount(listSourcesViaKernel(t, c), "minimax"); !ok {
			t.Fatal("开启渠道后应有 minimax 源")
		}
	})

	t.Run("备份导出含凭据", func(t *testing.T) {
		invoke(t, ps, "account_add", `{"platform":"cline","credential":"CL-KEY","nickname":"cl1"}`)
		exp := invoke(t, ps, "backup_export", `{}`)
		payload, _ := exp["payload"].(map[string]interface{})
		if payload == nil {
			t.Fatal("backup_export 应返回 payload")
		}
		creds, _ := payload["credentials"].(map[string]interface{})
		found := false
		for _, v := range creds {
			if v == "CL-KEY" {
				found = true
			}
		}
		if !found {
			t.Fatalf("备份应含凭据明文: %v", creds)
		}
	})
}

// TestChannelsExtendedRPCBehavior：对齐参考剩余 RPC 面（credits.status/balances、
// account.reset/retest、model.setDisabledMany/clearDead、provider 排序、
// gateway/backup.status、usage.badgePreference、onboarding）。
func TestChannelsExtendedRPCBehavior(t *testing.T) {
	_, ps := poolCore(t)

	t.Run("credits_balances_无凭据报错_有凭据解析", func(t *testing.T) {
		m := invoke(t, ps, "account_add", `{"platform":"minimax","nickname":"b1"}`)
		aid, _ := m["account_id"].(string)
		j := invoke(t, ps, "credits_balances", `{"platform":"minimax"}`)
		var row map[string]interface{}
		for _, a := range asArr(j["accounts"]) {
			x := a.(map[string]interface{})
			if x["accountId"] == aid {
				row = x
			}
		}
		if row == nil {
			t.Fatal("balances 应含该账号")
		}
		if row["balance"] != nil || row["error"] == nil {
			t.Fatalf("无凭据应 error 而非编造余额: %v", row)
		}
	})

	t.Run("credits_status_无端点平台如实null", func(t *testing.T) {
		m := invoke(t, ps, "account_add", `{"platform":"opencode","nickname":"oc1"}`)
		aid, _ := m["account_id"].(string)
		j := invoke(t, ps, "credits_status", `{"platform":"opencode"}`)
		for _, a := range asArr(j["accounts"]) {
			x := a.(map[string]interface{})
			if x["accountId"] == aid && x["status"] != nil {
				t.Fatalf("opencode 无状态端点应 null: %v", x)
			}
		}
	})

	t.Run("account_reset_clears", func(t *testing.T) {
		invoke(t, ps, "account_add", `{"platform":"minimax","nickname":"r1"}`)
		j := invoke(t, ps, "account_reset", `{"platform":"minimax","account_id":"__all__"}`)
		if j["ok"] != true {
			t.Fatalf("resetAll: %v", j)
		}
	})

	t.Run("model_set_disabled_many", func(t *testing.T) {
		invoke(t, ps, "model_set_disabled_many", `{"platform":"minimax","model_ids":["MiniMax-M3","MiniMax-M2.7-highspeed"],"disabled":true}`)
		j := invoke(t, ps, "model_list", `{"platform":"minimax"}`)
		off := 0
		for _, x := range asArr(j["models"]) {
			mm := x.(map[string]interface{})
			if mm["disabled"] == true {
				off++
			}
		}
		if off != 2 {
			t.Fatalf("批量关2模型后应2个disabled, 实得 %d", off)
		}
		invoke(t, ps, "model_set_disabled_many", `{"platform":"minimax","model_ids":["MiniMax-M3","MiniMax-M2.7-highspeed"],"disabled":false}`)
	})

	t.Run("model_clear_dead_只清黑名单里已下线的", func(t *testing.T) {
		invoke(t, ps, "model_set_disabled", `{"platform":"minimax","model_id":"MiniMax-M3","disabled":true}`)
		j := invoke(t, ps, "model_clear_dead", `{"platform":"minimax"}`)
		if n := len(asArr(j["removed"])); n != 0 {
			t.Fatalf("在册模型不算死模型, 实得 %d", n)
		}
		// 手动放一个不在 models 表里的黑名单项
		invoke(t, ps, "model_set_disabled_many", `{"platform":"minimax","model_ids":["ghost-model"],"disabled":true}`)
		j = invoke(t, ps, "model_clear_dead", `{"platform":"minimax"}`)
		if n := len(asArr(j["removed"])); n != 1 {
			t.Fatalf("死模型应被清: %v", j)
		}
		invoke(t, ps, "model_set_disabled", `{"platform":"minimax","model_id":"MiniMax-M3","disabled":false}`)
	})

	t.Run("provider_order", func(t *testing.T) {
		invoke(t, ps, "provider_set_order", `{"order":["minimax","trae","bogus"]}`)
		j := invoke(t, ps, "provider_get_order", `{}`)
		order := asArr(j["order"])
		if len(order) != 2 {
			t.Fatalf("bogus 应被过滤, 实得 %v", order)
		}
		if order[0] != "minimax" {
			t.Fatalf("顺序应保留: %v", order)
		}
	})

	t.Run("backup_status", func(t *testing.T) {
		invoke(t, ps, "account_add", `{"platform":"minimax","nickname":"s1"}`)
		j := invoke(t, ps, "backup_status", `{}`)
		if n, _ := j["accounts"].(float64); n < 1 {
			t.Fatalf("backup_status 应统计账号: %v", j)
		}
	})

	t.Run("badge_preference_非法值拒绝", func(t *testing.T) {
		j := invoke(t, ps, "usage_badge_preference", `{"preference":"bogus"}`)
		if j["error"] == nil {
			t.Fatalf("非法偏好应拒绝: %v", j)
		}
		j = invoke(t, ps, "usage_badge_preference", `{"preference":"credits"}`)
		if j["preference"] != "credits" {
			t.Fatalf("合法偏好应保存: %v", j)
		}
	})

	t.Run("onboarding_不支持平台如实报", func(t *testing.T) {
		j := invoke(t, ps, "onboarding_status", `{"platform":"minimax"}`)
		if j["supported"] != false {
			t.Fatalf("minimax 无新人任务应 supported=false: %v", j)
		}
	})

	t.Run("permanent_lock_读写", func(t *testing.T) {
		j := invoke(t, ps, "credits_permanent_lock", `{"platform":"buddy","locked":true}`)
		if j["locked"] != true {
			t.Fatalf("锁定应生效: %v", j)
		}
		j = invoke(t, ps, "credits_permanent_lock", `{"platform":"buddy"}`)
		if j["locked"] != true {
			t.Fatalf("只读应回显锁定: %v", j)
		}
	})
}

// TestBuddyBillingHeadersMatchReference 锁定 buddy 计费请求的头（逐行移植参考
// checkinHeaders）。真实缺陷：实现只发 Authorization，上游直接 401 —— 修好后
// 账号能读到凭据但余额仍失败，运维看到的仍是「拿不到账号信息」。
func TestBuddyBillingHeadersMatchReference(t *testing.T) {
	_, ps := poolCore(t)

	var got map[string]string
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = map[string]string{}
		for k := range r.Header {
			got[k] = r.Header.Get(k)
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"data":{"totalCredit":42}}`)
	}))
	defer ts.Close()

	// 把 buddy 的余额/签到端点指到本地假上游
	invoke(t, ps, "set_endpoints", fmt.Sprintf(
		`{"platform":"buddy","balance":{"url":%q,"method":"POST","body":"{}"}}`, ts.URL))
	m := invoke(t, ps, "account_add", `{"platform":"buddy","credential":"BD-AT","nickname":"b"}`)
	if m["ok"] != true {
		t.Fatalf("account_add: %v", m)
	}
	invoke(t, ps, "credits_balances", `{"platform":"buddy"}`)

	if got == nil {
		t.Fatal("余额请求未发出")
	}
	// 参考 checkinHeaders 的必备身份头（缺一个上游就 401）
	want := map[string]string{
		"Authorization":  "Bearer BD-AT",
		"X-Domain":       "copilot.tencent.com",
		"X-Product":      "SaaS",
		"X-Product-Code": "codebuddy",
		"User-Agent":     "CodeBuddyIDE/1.106.1",
	}
	for k, v := range want {
		if got[k] != v {
			t.Errorf("头 %s = %q, 期望 %q（buddy 计费端点缺它会被上游 401）", k, got[k], v)
		}
	}
}

// TestBuddyBalanceParseDoubleNested 逐行锁定参考 fetchCreditBalance 的解析口径：
// 响应是**双层嵌套** data.Response.Data.Accounts[]（与签到端点的单层 data 不同，
// 这是该接口最容易错的地方），且只累加**有效包**的 CycleCapacityRemain ——
// 失效包服务端仍返回，并进总额会让数字虚高（参考实测 155.67 → 655.67）。
func TestBuddyBalanceParseDoubleNested(t *testing.T) {
	_, ps := poolCore(t)

	// 两个有效包 + 一个失效包（ExpiredTime 在 2020 年）
	body := `{"code":0,"msg":"OK","data":{"Response":{"Data":{"Accounts":[
	  {"PackageName":"A","Status":1,"CycleCapacityRemain":100.5,"ExpiredTime":"2099-01-01 00:00:00"},
	  {"PackageName":"B","Status":1,"CycleCapacityRemain":55.17000031,"ExpiredTime":"2099-01-01 00:00:00"},
	  {"PackageName":"C","Status":1,"CycleCapacityRemain":500,"ExpiredTime":"2020-01-01 00:00:00"}
	]}}}}`
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, body)
	}))
	defer ts.Close()

	invoke(t, ps, "set_endpoints", fmt.Sprintf(
		`{"platform":"buddy","balance":{"url":%q,"method":"POST","body":"{}"}}`, ts.URL))
	invoke(t, ps, "account_add", `{"platform":"buddy","credential":"BD-AT","nickname":"b"}`)

	j := invoke(t, ps, "credits_balances", `{"platform":"buddy"}`)
	var row map[string]interface{}
	for _, a := range asArr(j["accounts"]) {
		row = a.(map[string]interface{})
	}
	bal, _ := row["balance"].(map[string]interface{})
	if bal == nil {
		t.Fatalf("应解析出余额, 实得 error=%v", row["error"])
	}
	// 100.5 + 55.17 = 155.67（失效包的 500 不计入）
	if got, _ := bal["total"].(float64); got != 155.67 {
		t.Fatalf("余额 = %v, 期望 155.67（只累加有效包、两位小数）", bal["total"])
	}

	// 单层 data（缺 Response 层）必须判为解析失败，不能当 0
	body = `{"code":0,"data":{"TotalDosage":609}}`
	j = invoke(t, ps, "credits_balances", `{"platform":"buddy"}`)
	for _, a := range asArr(j["accounts"]) {
		row = a.(map[string]interface{})
	}
	if row["balance"] != nil {
		t.Fatalf("双层缺失时应报错而非当 0（参考语义：区分「查不到」与「余额为 0」）: %v", row)
	}
}