| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229 |
- package service
- import (
- "context"
- "testing"
- "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
- "github.com/stretchr/testify/require"
- )
- type compositeOwnershipAccountRepo struct {
- AccountRepository
- accounts []Account
- }
- func (r *compositeOwnershipAccountRepo) ListSchedulableByGroupID(context.Context, int64) ([]Account, error) {
- return r.accounts, nil
- }
- // Scenario: 唯一平台的精确别名可路由
- func TestResolveCompositeModelOwnershipKeepsProviderAccountsIsolated(t *testing.T) {
- groupID := int64(7)
- repo := &compositeOwnershipAccountRepo{
- accounts: []Account{
- {
- ID: 1,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{
- "model_mapping": map[string]any{"gpt-public": "gpt-5"},
- },
- },
- {
- ID: 2,
- Platform: PlatformDeepseek,
- Credentials: map[string]any{
- "model_mapping": map[string]any{"reasoning-alias": "deepseek-v4-pro"},
- },
- },
- },
- }
- svc := &GatewayService{accountRepo: repo}
- deepSeekOwnership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "reasoning-alias")
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformDeepseek, Matched: true}, deepSeekOwnership)
- openAIOwnership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "gpt-public")
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, openAIOwnership)
- }
- // Scenario: 通配符和空映射不声明所有权
- func TestResolveCompositeModelOwnershipRequiresNonEmptyExactMappings(t *testing.T) {
- groupID := int64(7)
- repo := &compositeOwnershipAccountRepo{
- accounts: []Account{
- {
- ID: 1,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{
- "model_mapping": map[string]any{"*": "gpt-5", "gpt-*": "gpt-5", "empty-alias": ""},
- },
- },
- {
- ID: 2,
- Platform: PlatformGrok,
- Credentials: map[string]any{
- "model_mapping": map[string]any{"grok-public": "grok-4"},
- },
- },
- },
- }
- svc := &GatewayService{accountRepo: repo}
- for _, model := range []string{"gpt-5", "empty-alias", "unknown-alias"} {
- ownership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, model)
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{}, ownership, "model=%s", model)
- }
- ownership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "grok-public")
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformGrok, Matched: true}, ownership)
- }
- func TestResolveCompositeModelOwnershipAllowsSamePlatformAndRejectsCrossPlatformAliases(t *testing.T) {
- groupID := int64(7)
- repo := &compositeOwnershipAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Credentials: map[string]any{"model_mapping": map[string]any{"shared-openai": "gpt-5", "ambiguous": "gpt-5"}}},
- {ID: 2, Platform: PlatformOpenAI, Credentials: map[string]any{"model_mapping": map[string]any{"shared-openai": "gpt-5.1"}}},
- {ID: 3, Platform: PlatformDeepseek, Credentials: map[string]any{"model_mapping": map[string]any{"ambiguous": "deepseek-v4-pro"}}},
- },
- }
- svc := &GatewayService{accountRepo: repo}
- samePlatform, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "shared-openai")
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, samePlatform)
- ambiguous, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "ambiguous")
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{Ambiguous: true}, ambiguous)
- }
- func TestNewGatewayServiceWiresCompositeModelOwnershipResolver(t *testing.T) {
- groupID := int64(7)
- repo := &compositeOwnershipAccountRepo{
- accounts: []Account{{
- ID: 1,
- Platform: PlatformDeepseek,
- Credentials: map[string]any{"model_mapping": map[string]any{"reasoning-alias": "deepseek-v4-pro"}},
- }},
- }
- resolver := NewCompositeRouteResolver(nil)
- svc := NewGatewayService(
- repo,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- nil,
- resolver,
- nil,
- nil,
- )
- require.Same(t, resolver, svc.compositeResolver)
- decision, err := resolver.Resolve(context.Background(), groupID, "reasoning-alias", CompositeRouteEndpointResponses)
- require.NoError(t, err)
- require.True(t, decision.Matched)
- require.Equal(t, CompositeRouteSourceAccount, decision.Source)
- require.Equal(t, PlatformDeepseek, decision.TargetPlatform)
- }
- func TestDetectModelPlatform(t *testing.T) {
- tests := []struct {
- name string
- model string
- platform string
- ok bool
- }{
- {name: "claude", model: "claude-sonnet-4-5", platform: PlatformAnthropic, ok: true},
- {name: "anthropic prefix", model: "anthropic/claude-opus-4-5", platform: PlatformAnthropic, ok: true},
- {name: "gpt", model: "gpt-5.1", platform: PlatformOpenAI, ok: true},
- {name: "o series", model: "o3-mini", platform: PlatformOpenAI, ok: true},
- {name: "embedding", model: "text-embedding-3-large", platform: PlatformOpenAI, ok: true},
- {name: "gemini", model: "gemini-3-pro", platform: PlatformGemini, ok: true},
- {name: "gemini models prefix", model: "models/gemini-2.5-flash", platform: PlatformGemini, ok: true},
- {name: "learnlm", model: "learnlm-2.0-flash-experimental", platform: PlatformGemini, ok: true},
- {name: "grok", model: "grok-4", platform: PlatformGrok, ok: true},
- {name: "xai prefix", model: "xai/grok-4", platform: PlatformGrok, ok: true},
- {name: "kimi", model: "kimi-k2-thinking", platform: PlatformKimi, ok: true},
- {name: "kimi code bare k3", model: "K3", platform: PlatformKimi, ok: true},
- {name: "kimi code bare k3 256k", model: "k3-256k", platform: PlatformKimi, ok: true},
- {name: "kimi code provider prefix", model: "kimi-code/k3", platform: PlatformKimi, ok: true},
- {name: "moonshot prefix", model: "moonshot/moonshot-v1-32k", platform: PlatformKimi, ok: true},
- {name: "zhipu", model: "glm-5.2", platform: PlatformZhipu, ok: true},
- {name: "deepseek", model: "deepseek-v4-pro", platform: PlatformDeepseek, ok: true},
- {name: "minimax", model: "MiniMax-M3", platform: PlatformMiniMax, ok: true},
- {name: "minimax prefix", model: "minimax/MiniMax-M2.5", platform: PlatformMiniMax, ok: true},
- {name: "abab legacy", model: "abab6.5-chat", platform: PlatformMiniMax, ok: true},
- {name: "abab7 legacy", model: "abab7-chat-preview", platform: PlatformMiniMax, ok: true},
- {name: "abab unrelated namespace", model: "abab-other", ok: false},
- {name: "unknown k3 alias", model: "k3-preview", ok: false},
- {name: "unknown", model: "llama-4-maverick", ok: false},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- platform, ok := DetectModelPlatform(tt.model)
- require.Equal(t, tt.ok, ok)
- require.Equal(t, tt.platform, platform)
- })
- }
- }
- func TestQuotaPlatformCompositeUsesResolvedOrForceOnly(t *testing.T) {
- apiKey := &APIKey{Group: &Group{Platform: PlatformComposite}}
- require.Equal(t, "", QuotaPlatform(context.Background(), apiKey))
- require.Equal(t, PlatformGemini, QuotaPlatform(WithResolvedTargetPlatform(context.Background(), PlatformGemini), apiKey))
- require.Equal(t, PlatformAntigravity, QuotaPlatform(context.WithValue(context.Background(), ctxkey.ForcePlatform, PlatformAntigravity), apiKey))
- ctx := WithResolvedTargetPlatform(context.Background(), PlatformAnthropic)
- ctx = context.WithValue(ctx, ctxkey.ForcePlatform, PlatformAntigravity)
- require.Equal(t, PlatformAntigravity, QuotaPlatform(ctx, apiKey))
- }
- func TestCompositeGroupSchedulerHasAllCanonicalPlatformBuckets(t *testing.T) {
- seen := make(map[string]struct{})
- for _, bucket := range schedulerCanonicalBuckets(99) {
- seen[bucket.Platform] = struct{}{}
- }
- platforms := make([]string, 0, len(seen))
- for platform := range seen {
- platforms = append(platforms, platform)
- }
- require.ElementsMatch(t,
- []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax},
- platforms,
- )
- }
- func TestCompositeConcretePlatformsIncludeCNProviders(t *testing.T) {
- for _, platform := range []string{PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax} {
- require.True(t, isConcreteRequestPlatform(platform))
- require.True(t, canCopyAccountsFromGroupPlatform(PlatformComposite, platform))
- }
- }
|