sub2api-cn-relay-manager/internal/provision/model_pool_test.go

package provision

import (
	"strings"
	"testing"

	"sub2api-cn-relay-manager/internal/host/sub2api"
	"sub2api-cn-relay-manager/internal/pack"
)

func TestBuildModelPoolFiltersUnsupportedRoutesAndSortsByPriority(t *testing.T) {
	t.Parallel()

	providerA := sampleProviderManifest()
	providerA.ProviderID = "deepseek-official"
	providerA.DisplayName = "DeepSeek Official"
	providerA.BaseURL = "https://api.deepseek.com/v1"

	providerB := sampleProviderManifest()
	providerB.ProviderID = "deepseek-backup"
	providerB.DisplayName = "DeepSeek Backup"
	providerB.BaseURL = "https://backup.deepseek.example.com/v1"

	providerC := sampleProviderManifest()
	providerC.ProviderID = "deepseek-bad"
	providerC.DisplayName = "DeepSeek Bad"
	providerC.BaseURL = "https://bad.deepseek.example.com/v1"

	pool, err := BuildModelPool(ModelPoolBuildRequest{
		PublicModel: "deepseek-chat",
		Candidates: []ModelPoolCandidate{
			{
				RouteID:   "route-backup",
				Provider:  providerB,
				Priority:  20,
				Inventory: capabilityInventoryWithSupport("deepseek-chat", "deepseek-chat", sub2api.SupportLevelDirect, true, true),
			},
			{
				RouteID:   "route-primary",
				Provider:  providerA,
				Priority:  10,
				Inventory: capabilityInventoryWithSupport("deepseek-chat", "deepseek-chat", sub2api.SupportLevelDirect, true, true),
			},
			{
				RouteID:   "route-bad",
				Provider:  providerC,
				Priority:  5,
				Inventory: capabilityInventoryWithSupport("deepseek-chat", "deepseek-chat", sub2api.SupportLevelUpstreamUnhealthy, true, false),
			},
		},
	})
	if err != nil {
		t.Fatalf("BuildModelPool() error = %v", err)
	}
	if len(pool.Routes) != 2 {
		t.Fatalf("len(pool.Routes) = %d, want 2", len(pool.Routes))
	}
	if pool.Routes[0].RouteID != "route-primary" || pool.Routes[1].RouteID != "route-backup" {
		t.Fatalf("route order = [%s %s], want [route-primary route-backup]", pool.Routes[0].RouteID, pool.Routes[1].RouteID)
	}
	if pool.Routes[0].SupportLevel != sub2api.SupportLevelDirect {
		t.Fatalf("pool.Routes[0].SupportLevel = %q, want %q", pool.Routes[0].SupportLevel, sub2api.SupportLevelDirect)
	}
}

func TestBuildModelPoolPreservesAdvertisedAndCallableModelDifference(t *testing.T) {
	t.Parallel()

	provider := sampleProviderManifest()
	provider.ProviderID = "deepseek-alias"
	provider.DisplayName = "DeepSeek Alias"
	provider.BaseURL = "https://alias.deepseek.example.com/v1"

	pool, err := BuildModelPool(ModelPoolBuildRequest{
		PublicModel: "deepseek-chat",
		Candidates: []ModelPoolCandidate{
			{
				RouteID:         "route-alias",
				Provider:        provider,
				Priority:        10,
				AdvertisedModel: "deepseek-v3",
				CallableModel:   "deepseek-chat",
				Inventory:       capabilityInventoryWithSupport("deepseek-chat", "deepseek-chat", sub2api.SupportLevelDirect, true, true),
			},
		},
	})
	if err != nil {
		t.Fatalf("BuildModelPool() error = %v", err)
	}
	if len(pool.Routes) != 1 {
		t.Fatalf("len(pool.Routes) = %d, want 1", len(pool.Routes))
	}
	if pool.Routes[0].AdvertisedModel != "deepseek-v3" {
		t.Fatalf("AdvertisedModel = %q, want deepseek-v3", pool.Routes[0].AdvertisedModel)
	}
	if pool.Routes[0].CallableModel != "deepseek-chat" {
		t.Fatalf("CallableModel = %q, want deepseek-chat", pool.Routes[0].CallableModel)
	}
}

func TestBuildModelPoolAllowsChatOnlyRouteWhenExplicitlyEnabled(t *testing.T) {
	t.Parallel()

	pool, err := BuildModelPool(ModelPoolBuildRequest{
		PublicModel:                  "kimi-k2.6",
		AllowPluginAdapterCandidates: true,
		Candidates: []ModelPoolCandidate{
			{
				RouteID:   "route-kimi",
				Provider:  pack.ProviderManifest{ProviderID: "kimi-a7m", DisplayName: "Kimi", BaseURL: "https://kimi.a7m.com.cn/v1"},
				Priority:  10,
				Inventory: capabilityInventoryWithSupport("kimi-k2.6", "kimi-k2.6", sub2api.SupportLevelWithPluginAdapter, true, false),
			},
		},
	})
	if err != nil {
		t.Fatalf("BuildModelPool() error = %v", err)
	}
	if len(pool.Routes) != 1 {
		t.Fatalf("len(pool.Routes) = %d, want 1", len(pool.Routes))
	}
	if pool.Routes[0].SupportLevel != sub2api.SupportLevelWithPluginAdapter {
		t.Fatalf("SupportLevel = %q, want %q", pool.Routes[0].SupportLevel, sub2api.SupportLevelWithPluginAdapter)
	}
	if pool.Routes[0].SupportsResponses {
		t.Fatal("SupportsResponses = true, want false")
	}
}

func TestBuildModelPoolReturnsErrorWhenNoEligibleRoutesRemain(t *testing.T) {
	t.Parallel()

	_, err := BuildModelPool(ModelPoolBuildRequest{
		PublicModel: "kimi-k2.6",
		Candidates: []ModelPoolCandidate{{
			RouteID:   "route-kimi",
			Provider:  pack.ProviderManifest{ProviderID: "kimi-a7m", DisplayName: "Kimi", BaseURL: "https://kimi.a7m.com.cn/v1"},
			Priority:  10,
			Inventory: capabilityInventoryWithSupport("kimi-k2.6", "kimi-k2.6", sub2api.SupportLevelWithPluginAdapter, true, false),
		}},
	})
	if err == nil || !strings.Contains(err.Error(), "no eligible routes") {
		t.Fatalf("BuildModelPool() error = %v, want no eligible routes", err)
	}
}

func TestBuildModelPoolFiltersHostNotReadyRoute(t *testing.T) {
	t.Parallel()

	_, err := BuildModelPool(ModelPoolBuildRequest{
		PublicModel: "deepseek-chat",
		Candidates: []ModelPoolCandidate{{
			RouteID:   "route-host-not-ready",
			Provider:  sampleProviderManifest(),
			Priority:  10,
			Inventory: capabilityInventoryWithHostReady("deepseek-chat", "deepseek-chat", sub2api.SupportLevelDirect, true, true, false),
		}},
	})
	if err == nil || !strings.Contains(err.Error(), "no eligible routes") {
		t.Fatalf("BuildModelPool() error = %v, want no eligible routes after host_ready filter", err)
	}
}

func TestBuildModelPoolFiltersUnschedulableRoute(t *testing.T) {
	t.Parallel()

	unschedulable := false
	_, err := BuildModelPool(ModelPoolBuildRequest{
		PublicModel: "deepseek-chat",
		Candidates: []ModelPoolCandidate{{
			RouteID:     "route-unschedulable",
			Provider:    sampleProviderManifest(),
			Priority:    10,
			Schedulable: &unschedulable,
			Inventory:   capabilityInventoryWithSupport("deepseek-chat", "deepseek-chat", sub2api.SupportLevelDirect, true, true),
		}},
	})
	if err == nil || !strings.Contains(err.Error(), "no eligible routes") {
		t.Fatalf("BuildModelPool() error = %v, want no eligible routes after schedulable filter", err)
	}
}

func TestBuildModelPoolPreservesSupportedModelsForRoute(t *testing.T) {
	t.Parallel()

	pool, err := BuildModelPool(ModelPoolBuildRequest{
		PublicModel: "deepseek-chat",
		Candidates: []ModelPoolCandidate{{
			RouteID:  "route-primary",
			Provider: sampleProviderManifest(),
			Priority: 10,
			Inventory: capabilityInventoryWithMultiModels([]sub2api.ModelCapabilitySummary{
				{
					RawModelID:           "deepseek-chat",
					CanonicalModelFamily: "deepseek-chat",
					SmokeChatOK:          true,
					SupportLevel:         sub2api.SupportLevelDirect,
				},
				{
					RawModelID:           "deepseek-reasoner",
					CanonicalModelFamily: "deepseek-reasoner",
					SmokeChatOK:          true,
					SupportLevel:         sub2api.SupportLevelDirect,
				},
			}),
		}},
	})
	if err != nil {
		t.Fatalf("BuildModelPool() error = %v", err)
	}
	if len(pool.Routes) != 1 {
		t.Fatalf("len(pool.Routes) = %d, want 1", len(pool.Routes))
	}
	if len(pool.Routes[0].SupportedModels) != 2 {
		t.Fatalf("len(pool.Routes[0].SupportedModels) = %d, want 2", len(pool.Routes[0].SupportedModels))
	}
	if pool.Routes[0].SupportedModels[0] != "deepseek-chat" || pool.Routes[0].SupportedModels[1] != "deepseek-reasoner" {
		t.Fatalf("SupportedModels = %#v, want [deepseek-chat deepseek-reasoner]", pool.Routes[0].SupportedModels)
	}
}

func capabilityInventoryWithSupport(rawModel string, canonical string, supportLevel string, chatOK bool, responsesOK bool) sub2api.CapabilityInventory {
	return capabilityInventoryWithHostReady(rawModel, canonical, supportLevel, chatOK, responsesOK, true)
}

func capabilityInventoryWithHostReady(rawModel string, canonical string, supportLevel string, chatOK bool, responsesOK bool, hostReady bool) sub2api.CapabilityInventory {
	return capabilityInventoryWithMultiModelsHostReady([]sub2api.ModelCapabilitySummary{{
		RawModelID:           rawModel,
		CanonicalModelFamily: canonical,
		SmokeChatOK:          chatOK,
		SupportLevel:         supportLevel,
		KnownAdvisories: func() []string {
			if responsesOK {
				return nil
			}
			return []string{"responses_unsupported_but_chat_ok"}
		}(),
	}}, hostReady)
}

func capabilityInventoryWithMultiModels(models []sub2api.ModelCapabilitySummary) sub2api.CapabilityInventory {
	return capabilityInventoryWithMultiModelsHostReady(models, true)
}

func capabilityInventoryWithMultiModelsHostReady(models []sub2api.ModelCapabilitySummary, hostReady bool) sub2api.CapabilityInventory {
	return sub2api.CapabilityInventory{
		HostReady: hostReady,
		Host:      sub2api.HostCapabilities{Groups: true, Channels: true, Accounts: true, AccountTest: true},
		Models:    models,
	}
}