300 lines
9.1 KiB
Go
300 lines
9.1 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"testing"
|
||
)
|
||
|
||
// 目录与选择分两张表，是为了让「已选」在模型从平台目录里消失后仍然留存。
|
||
// 这个测试钉住那个行为 —— 合并成一张带 allowed 标记的表就会失败。
|
||
func TestAllowedModelsSurviveCatalogChurn(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||
{Provider: "llmsproxy", Model: "AUTO", DisplayName: "AUTO"},
|
||
{Provider: "llmsproxy", Model: "claude-sonnet-4-6"},
|
||
}); err != nil {
|
||
t.Fatalf("首次上报目录: %v", err)
|
||
}
|
||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{
|
||
{Provider: "llmsproxy", Model: "AUTO"},
|
||
}); err != nil {
|
||
t.Fatalf("保存选择: %v", err)
|
||
}
|
||
|
||
// 平台侧 AUTO 临时下线，只上报另一个
|
||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||
{Provider: "llmsproxy", Model: "claude-sonnet-4-6"},
|
||
}); err != nil {
|
||
t.Fatalf("二次上报目录: %v", err)
|
||
}
|
||
|
||
allowed, err := ListAllowedModels(ctx, "dsh")
|
||
if err != nil {
|
||
t.Fatalf("ListAllowedModels: %v", err)
|
||
}
|
||
if len(allowed) != 1 || allowed[0].Model != "AUTO" {
|
||
t.Fatalf("模型从目录消失后选择也被删了：%+v —— "+
|
||
"两张表分开存的意义就在于此", allowed)
|
||
}
|
||
|
||
// 它应当被标为 stale，界面上才能提示「已选但平台没上报」
|
||
stale, err := ListStaleAllowedModels(ctx, "dsh")
|
||
if err != nil {
|
||
t.Fatalf("ListStaleAllowedModels: %v", err)
|
||
}
|
||
if len(stale) != 1 || stale[0].Model != "AUTO" {
|
||
t.Errorf("应有 1 个 stale，实际 %+v", stale)
|
||
}
|
||
|
||
// 模型回来后不该再是 stale，也不需要重新勾选
|
||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||
{Provider: "llmsproxy", Model: "AUTO"},
|
||
{Provider: "llmsproxy", Model: "claude-sonnet-4-6"},
|
||
}); err != nil {
|
||
t.Fatalf("三次上报: %v", err)
|
||
}
|
||
stale2, _ := ListStaleAllowedModels(ctx, "dsh")
|
||
if len(stale2) != 0 {
|
||
t.Errorf("模型回来后不该再是 stale：%+v", stale2)
|
||
}
|
||
}
|
||
|
||
// 目录整表替换：平台下线的模型必须从配置页消失，
|
||
// 否则管理员会勾选一个平台其实调不到的模型。
|
||
func TestReplaceModelCatalogIsFullReplace(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
if err := ReplaceModelCatalog(ctx, "opencode", []CatalogModel{
|
||
{Provider: "p", Model: "a"},
|
||
{Provider: "p", Model: "b"},
|
||
}); err != nil {
|
||
t.Fatalf("首次: %v", err)
|
||
}
|
||
if err := ReplaceModelCatalog(ctx, "opencode", []CatalogModel{
|
||
{Provider: "p", Model: "a"},
|
||
}); err != nil {
|
||
t.Fatalf("二次: %v", err)
|
||
}
|
||
got, err := ListModelCatalog(ctx, "opencode")
|
||
if err != nil {
|
||
t.Fatalf("ListModelCatalog: %v", err)
|
||
}
|
||
if len(got) != 1 || got[0].Model != "a" {
|
||
t.Fatalf("整表替换失效：%+v", got)
|
||
}
|
||
}
|
||
|
||
// ListModelCatalog 要在同一次查询里标出「已选」与 rank，
|
||
// 前端才能画出带勾选与顺序的清单。
|
||
func TestListModelCatalogMarksAllowedAndRank(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||
{Provider: "p", Model: "first", DisplayName: "第一"},
|
||
{Provider: "p", Model: "second"},
|
||
{Provider: "p", Model: "unpicked"},
|
||
}); err != nil {
|
||
t.Fatalf("上报: %v", err)
|
||
}
|
||
// 顺序即优先级：second 排前面
|
||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{
|
||
{Provider: "p", Model: "second"},
|
||
{Provider: "p", Model: "first"},
|
||
}); err != nil {
|
||
t.Fatalf("保存: %v", err)
|
||
}
|
||
|
||
got, err := ListModelCatalog(ctx, "dsh")
|
||
if err != nil {
|
||
t.Fatalf("ListModelCatalog: %v", err)
|
||
}
|
||
byModel := map[string]CatalogModel{}
|
||
for _, m := range got {
|
||
byModel[m.Model] = m
|
||
}
|
||
if !byModel["second"].Allowed || byModel["second"].Rank != 0 {
|
||
t.Errorf("second 应为 rank 0 的已选项：%+v", byModel["second"])
|
||
}
|
||
if !byModel["first"].Allowed || byModel["first"].Rank != 1 {
|
||
t.Errorf("first 应为 rank 1 的已选项：%+v", byModel["first"])
|
||
}
|
||
if byModel["unpicked"].Allowed {
|
||
t.Error("unpicked 不该被标为已选")
|
||
}
|
||
if byModel["first"].DisplayName != "第一" {
|
||
t.Errorf("display_name 未带出：%q", byModel["first"].DisplayName)
|
||
}
|
||
}
|
||
|
||
// 顺序就是插件的降级顺序，必须原样保存。
|
||
func TestSetAllowedModelsPreservesOrder(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
want := []ModelRef{
|
||
{Provider: "c", Model: "3"},
|
||
{Provider: "a", Model: "1"},
|
||
{Provider: "b", Model: "2"},
|
||
}
|
||
if err := SetAllowedModels(ctx, "dsh", want); err != nil {
|
||
t.Fatalf("SetAllowedModels: %v", err)
|
||
}
|
||
got, err := ListAllowedModels(ctx, "dsh")
|
||
if err != nil {
|
||
t.Fatalf("ListAllowedModels: %v", err)
|
||
}
|
||
if len(got) != len(want) {
|
||
t.Fatalf("数量不符：%d vs %d", len(got), len(want))
|
||
}
|
||
for i := range want {
|
||
if got[i] != want[i] {
|
||
t.Fatalf("第 %d 项顺序错：%+v，期望 %+v —— "+
|
||
"顺序就是插件的降级顺序，不能按字典序重排", i, got[i], want[i])
|
||
}
|
||
}
|
||
}
|
||
|
||
// 空列表表示「不限定」，是合法输入。
|
||
// 报错会让「取消所有限定」变成一件做不到的事。
|
||
func TestSetAllowedModelsAcceptsEmpty(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{{Provider: "p", Model: "m"}}); err != nil {
|
||
t.Fatalf("先设一个: %v", err)
|
||
}
|
||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{}); err != nil {
|
||
t.Fatalf("清空应当合法: %v", err)
|
||
}
|
||
got, _ := ListAllowedModels(ctx, "dsh")
|
||
if len(got) != 0 {
|
||
t.Errorf("清空后应为空，实际 %+v", got)
|
||
}
|
||
}
|
||
|
||
// 重复项跳过而不报错：它对最终顺序没有影响，
|
||
// 为一次无害的重复让整次保存失败只会让人以为配置没生效。
|
||
func TestSetAllowedModelsSkipsDuplicates(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
err := SetAllowedModels(ctx, "dsh", []ModelRef{
|
||
{Provider: "p", Model: "m"},
|
||
{Provider: "p", Model: "m"},
|
||
{Provider: "p", Model: "other"},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("重复项不该报错: %v", err)
|
||
}
|
||
got, _ := ListAllowedModels(ctx, "dsh")
|
||
if len(got) != 2 {
|
||
t.Fatalf("应保留 2 项，实际 %+v", got)
|
||
}
|
||
// rank 要连续：跳过重复项后不该在序号上留空洞
|
||
if got[0].Model != "m" || got[1].Model != "other" {
|
||
t.Errorf("顺序错：%+v", got)
|
||
}
|
||
}
|
||
|
||
// 字段不全的项跳过：半条记录在配置页上是一个点不动的空复选框。
|
||
func TestModelCatalogSkipsIncomplete(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||
{Provider: "", Model: "m"},
|
||
{Provider: "p", Model: ""},
|
||
{Provider: " ", Model: " "},
|
||
{Provider: "p", Model: "ok"},
|
||
}); err != nil {
|
||
t.Fatalf("上报: %v", err)
|
||
}
|
||
got, _ := ListModelCatalog(ctx, "dsh")
|
||
if len(got) != 1 || got[0].Model != "ok" {
|
||
t.Fatalf("应只留 1 项：%+v", got)
|
||
}
|
||
}
|
||
|
||
// 目录里重复的 provider/model 不该让整次事务失败（主键冲突）。
|
||
func TestReplaceModelCatalogDedupes(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||
{Provider: "p", Model: "m", DisplayName: "第一次"},
|
||
{Provider: "p", Model: "m", DisplayName: "第二次"},
|
||
}); err != nil {
|
||
t.Fatalf("重复不该报错: %v", err)
|
||
}
|
||
got, _ := ListModelCatalog(ctx, "dsh")
|
||
if len(got) != 1 {
|
||
t.Fatalf("应去重到 1 项：%+v", got)
|
||
}
|
||
if got[0].DisplayName != "第一次" {
|
||
t.Errorf("应保留第一条：%q", got[0].DisplayName)
|
||
}
|
||
}
|
||
|
||
// 各 Agent 的目录与选择互不影响。
|
||
func TestModelScopeIsolatedPerAgent(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{{Provider: "p", Model: "dsh-only"}}); err != nil {
|
||
t.Fatalf("dsh 上报: %v", err)
|
||
}
|
||
if err := ReplaceModelCatalog(ctx, "opencode", []CatalogModel{{Provider: "p", Model: "oc-only"}}); err != nil {
|
||
t.Fatalf("opencode 上报: %v", err)
|
||
}
|
||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{{Provider: "p", Model: "dsh-only"}}); err != nil {
|
||
t.Fatalf("dsh 选择: %v", err)
|
||
}
|
||
|
||
ocCatalog, _ := ListModelCatalog(ctx, "opencode")
|
||
if len(ocCatalog) != 1 || ocCatalog[0].Model != "oc-only" {
|
||
t.Fatalf("opencode 的目录被污染：%+v", ocCatalog)
|
||
}
|
||
if ocCatalog[0].Allowed {
|
||
t.Error("dsh 的选择串到 opencode 上了")
|
||
}
|
||
ocAllowed, _ := ListAllowedModels(ctx, "opencode")
|
||
if len(ocAllowed) != 0 {
|
||
t.Errorf("opencode 不该有已选项：%+v", ocAllowed)
|
||
}
|
||
}
|
||
|
||
// 上报数量超上限时截断而不报错：平台把上游几千个模型全列出来是它的自由，
|
||
// 但配置页上几千个复选框对人没有用。
|
||
func TestReplaceModelCatalogCaps(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
many := make([]CatalogModel, maxCatalogModels+50)
|
||
for i := range many {
|
||
many[i] = CatalogModel{Provider: "p", Model: string(rune('a'+i%26)) + itoaTest(i)}
|
||
}
|
||
if err := ReplaceModelCatalog(ctx, "dsh", many); err != nil {
|
||
t.Fatalf("超量上报不该报错: %v", err)
|
||
}
|
||
got, _ := ListModelCatalog(ctx, "dsh")
|
||
if len(got) != maxCatalogModels {
|
||
t.Errorf("应截断到 %d，实际 %d", maxCatalogModels, len(got))
|
||
}
|
||
}
|
||
|
||
func itoaTest(n int) string {
|
||
if n == 0 {
|
||
return "0"
|
||
}
|
||
var b []byte
|
||
for n > 0 {
|
||
b = append([]byte{byte('0' + n%10)}, b...)
|
||
n /= 10
|
||
}
|
||
return string(b)
|
||
}
|