composite_platform_test.go 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229
  1. package service
  2. import (
  3. "context"
  4. "testing"
  5. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  6. "github.com/stretchr/testify/require"
  7. )
  8. type compositeOwnershipAccountRepo struct {
  9. AccountRepository
  10. accounts []Account
  11. }
  12. func (r *compositeOwnershipAccountRepo) ListSchedulableByGroupID(context.Context, int64) ([]Account, error) {
  13. return r.accounts, nil
  14. }
  15. // Scenario: 唯一平台的精确别名可路由
  16. func TestResolveCompositeModelOwnershipKeepsProviderAccountsIsolated(t *testing.T) {
  17. groupID := int64(7)
  18. repo := &compositeOwnershipAccountRepo{
  19. accounts: []Account{
  20. {
  21. ID: 1,
  22. Platform: PlatformOpenAI,
  23. Credentials: map[string]any{
  24. "model_mapping": map[string]any{"gpt-public": "gpt-5"},
  25. },
  26. },
  27. {
  28. ID: 2,
  29. Platform: PlatformDeepseek,
  30. Credentials: map[string]any{
  31. "model_mapping": map[string]any{"reasoning-alias": "deepseek-v4-pro"},
  32. },
  33. },
  34. },
  35. }
  36. svc := &GatewayService{accountRepo: repo}
  37. deepSeekOwnership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "reasoning-alias")
  38. require.NoError(t, err)
  39. require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformDeepseek, Matched: true}, deepSeekOwnership)
  40. openAIOwnership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "gpt-public")
  41. require.NoError(t, err)
  42. require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, openAIOwnership)
  43. }
  44. // Scenario: 通配符和空映射不声明所有权
  45. func TestResolveCompositeModelOwnershipRequiresNonEmptyExactMappings(t *testing.T) {
  46. groupID := int64(7)
  47. repo := &compositeOwnershipAccountRepo{
  48. accounts: []Account{
  49. {
  50. ID: 1,
  51. Platform: PlatformOpenAI,
  52. Credentials: map[string]any{
  53. "model_mapping": map[string]any{"*": "gpt-5", "gpt-*": "gpt-5", "empty-alias": ""},
  54. },
  55. },
  56. {
  57. ID: 2,
  58. Platform: PlatformGrok,
  59. Credentials: map[string]any{
  60. "model_mapping": map[string]any{"grok-public": "grok-4"},
  61. },
  62. },
  63. },
  64. }
  65. svc := &GatewayService{accountRepo: repo}
  66. for _, model := range []string{"gpt-5", "empty-alias", "unknown-alias"} {
  67. ownership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, model)
  68. require.NoError(t, err)
  69. require.Equal(t, CompositeModelOwnership{}, ownership, "model=%s", model)
  70. }
  71. ownership, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "grok-public")
  72. require.NoError(t, err)
  73. require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformGrok, Matched: true}, ownership)
  74. }
  75. func TestResolveCompositeModelOwnershipAllowsSamePlatformAndRejectsCrossPlatformAliases(t *testing.T) {
  76. groupID := int64(7)
  77. repo := &compositeOwnershipAccountRepo{
  78. accounts: []Account{
  79. {ID: 1, Platform: PlatformOpenAI, Credentials: map[string]any{"model_mapping": map[string]any{"shared-openai": "gpt-5", "ambiguous": "gpt-5"}}},
  80. {ID: 2, Platform: PlatformOpenAI, Credentials: map[string]any{"model_mapping": map[string]any{"shared-openai": "gpt-5.1"}}},
  81. {ID: 3, Platform: PlatformDeepseek, Credentials: map[string]any{"model_mapping": map[string]any{"ambiguous": "deepseek-v4-pro"}}},
  82. },
  83. }
  84. svc := &GatewayService{accountRepo: repo}
  85. samePlatform, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "shared-openai")
  86. require.NoError(t, err)
  87. require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, samePlatform)
  88. ambiguous, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "ambiguous")
  89. require.NoError(t, err)
  90. require.Equal(t, CompositeModelOwnership{Ambiguous: true}, ambiguous)
  91. }
  92. func TestNewGatewayServiceWiresCompositeModelOwnershipResolver(t *testing.T) {
  93. groupID := int64(7)
  94. repo := &compositeOwnershipAccountRepo{
  95. accounts: []Account{{
  96. ID: 1,
  97. Platform: PlatformDeepseek,
  98. Credentials: map[string]any{"model_mapping": map[string]any{"reasoning-alias": "deepseek-v4-pro"}},
  99. }},
  100. }
  101. resolver := NewCompositeRouteResolver(nil)
  102. svc := NewGatewayService(
  103. repo,
  104. nil,
  105. nil,
  106. nil,
  107. nil,
  108. nil,
  109. nil,
  110. nil,
  111. nil,
  112. nil,
  113. nil,
  114. nil,
  115. nil,
  116. nil,
  117. nil,
  118. nil,
  119. nil,
  120. nil,
  121. nil,
  122. nil,
  123. nil,
  124. nil,
  125. nil,
  126. nil,
  127. nil,
  128. resolver,
  129. nil,
  130. nil,
  131. )
  132. require.Same(t, resolver, svc.compositeResolver)
  133. decision, err := resolver.Resolve(context.Background(), groupID, "reasoning-alias", CompositeRouteEndpointResponses)
  134. require.NoError(t, err)
  135. require.True(t, decision.Matched)
  136. require.Equal(t, CompositeRouteSourceAccount, decision.Source)
  137. require.Equal(t, PlatformDeepseek, decision.TargetPlatform)
  138. }
  139. func TestDetectModelPlatform(t *testing.T) {
  140. tests := []struct {
  141. name string
  142. model string
  143. platform string
  144. ok bool
  145. }{
  146. {name: "claude", model: "claude-sonnet-4-5", platform: PlatformAnthropic, ok: true},
  147. {name: "anthropic prefix", model: "anthropic/claude-opus-4-5", platform: PlatformAnthropic, ok: true},
  148. {name: "gpt", model: "gpt-5.1", platform: PlatformOpenAI, ok: true},
  149. {name: "o series", model: "o3-mini", platform: PlatformOpenAI, ok: true},
  150. {name: "embedding", model: "text-embedding-3-large", platform: PlatformOpenAI, ok: true},
  151. {name: "gemini", model: "gemini-3-pro", platform: PlatformGemini, ok: true},
  152. {name: "gemini models prefix", model: "models/gemini-2.5-flash", platform: PlatformGemini, ok: true},
  153. {name: "learnlm", model: "learnlm-2.0-flash-experimental", platform: PlatformGemini, ok: true},
  154. {name: "grok", model: "grok-4", platform: PlatformGrok, ok: true},
  155. {name: "xai prefix", model: "xai/grok-4", platform: PlatformGrok, ok: true},
  156. {name: "kimi", model: "kimi-k2-thinking", platform: PlatformKimi, ok: true},
  157. {name: "kimi code bare k3", model: "K3", platform: PlatformKimi, ok: true},
  158. {name: "kimi code bare k3 256k", model: "k3-256k", platform: PlatformKimi, ok: true},
  159. {name: "kimi code provider prefix", model: "kimi-code/k3", platform: PlatformKimi, ok: true},
  160. {name: "moonshot prefix", model: "moonshot/moonshot-v1-32k", platform: PlatformKimi, ok: true},
  161. {name: "zhipu", model: "glm-5.2", platform: PlatformZhipu, ok: true},
  162. {name: "deepseek", model: "deepseek-v4-pro", platform: PlatformDeepseek, ok: true},
  163. {name: "minimax", model: "MiniMax-M3", platform: PlatformMiniMax, ok: true},
  164. {name: "minimax prefix", model: "minimax/MiniMax-M2.5", platform: PlatformMiniMax, ok: true},
  165. {name: "abab legacy", model: "abab6.5-chat", platform: PlatformMiniMax, ok: true},
  166. {name: "abab7 legacy", model: "abab7-chat-preview", platform: PlatformMiniMax, ok: true},
  167. {name: "abab unrelated namespace", model: "abab-other", ok: false},
  168. {name: "unknown k3 alias", model: "k3-preview", ok: false},
  169. {name: "unknown", model: "llama-4-maverick", ok: false},
  170. }
  171. for _, tt := range tests {
  172. t.Run(tt.name, func(t *testing.T) {
  173. platform, ok := DetectModelPlatform(tt.model)
  174. require.Equal(t, tt.ok, ok)
  175. require.Equal(t, tt.platform, platform)
  176. })
  177. }
  178. }
  179. func TestQuotaPlatformCompositeUsesResolvedOrForceOnly(t *testing.T) {
  180. apiKey := &APIKey{Group: &Group{Platform: PlatformComposite}}
  181. require.Equal(t, "", QuotaPlatform(context.Background(), apiKey))
  182. require.Equal(t, PlatformGemini, QuotaPlatform(WithResolvedTargetPlatform(context.Background(), PlatformGemini), apiKey))
  183. require.Equal(t, PlatformAntigravity, QuotaPlatform(context.WithValue(context.Background(), ctxkey.ForcePlatform, PlatformAntigravity), apiKey))
  184. ctx := WithResolvedTargetPlatform(context.Background(), PlatformAnthropic)
  185. ctx = context.WithValue(ctx, ctxkey.ForcePlatform, PlatformAntigravity)
  186. require.Equal(t, PlatformAntigravity, QuotaPlatform(ctx, apiKey))
  187. }
  188. func TestCompositeGroupSchedulerHasAllCanonicalPlatformBuckets(t *testing.T) {
  189. seen := make(map[string]struct{})
  190. for _, bucket := range schedulerCanonicalBuckets(99) {
  191. seen[bucket.Platform] = struct{}{}
  192. }
  193. platforms := make([]string, 0, len(seen))
  194. for platform := range seen {
  195. platforms = append(platforms, platform)
  196. }
  197. require.ElementsMatch(t,
  198. []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax},
  199. platforms,
  200. )
  201. }
  202. func TestCompositeConcretePlatformsIncludeCNProviders(t *testing.T) {
  203. for _, platform := range []string{PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax} {
  204. require.True(t, isConcreteRequestPlatform(platform))
  205. require.True(t, canCopyAccountsFromGroupPlatform(PlatformComposite, platform))
  206. }
  207. }