openai_account_scheduler_spark_route_test.go 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. package service
  2. import (
  3. "context"
  4. "testing"
  5. "time"
  6. "github.com/Wei-Shaw/sub2api/internal/config"
  7. "github.com/stretchr/testify/require"
  8. )
  9. func TestSparkRoutingByModel(t *testing.T) {
  10. ctx := context.Background()
  11. sparkModel := "gpt-5.3-codex-spark"
  12. normalModel := "gpt-5.3-codex"
  13. sparkCreds := map[string]any{"model_mapping": defaultSparkShadowModelMapping()}
  14. newScheduler := func(snapshot map[int64]*Account) *defaultOpenAIAccountScheduler {
  15. return &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{
  16. schedulerSnapshot: &SchedulerSnapshotService{
  17. cache: &openAISnapshotCacheStub{accountsByID: snapshot},
  18. },
  19. cfg: &config.Config{},
  20. }}
  21. }
  22. sparkReq := OpenAIAccountScheduleRequest{RequestedModel: sparkModel, Platform: PlatformOpenAI}
  23. normalReq := OpenAIAccountScheduleRequest{RequestedModel: normalModel, Platform: PlatformOpenAI}
  24. t.Run("normal_account_with_spark_mapping_accepts_spark", func(t *testing.T) {
  25. acc := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Credentials: sparkCreds}
  26. require.True(t, newScheduler(nil).isAccountRequestCompatible(ctx, acc, sparkReq),
  27. "普通账号配了 spark → 可承接 spark(类型门已移除)")
  28. })
  29. t.Run("normal_account_without_spark_rejects_spark", func(t *testing.T) {
  30. acc := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true,
  31. Credentials: map[string]any{"model_mapping": map[string]any{normalModel: normalModel}}}
  32. require.False(t, newScheduler(nil).isAccountRequestCompatible(ctx, acc, sparkReq),
  33. "普通账号未配 spark → 拒 spark(按配置而非类型)")
  34. })
  35. t.Run("shadow_with_spark_mapping_accepts_spark_rejects_non_spark", func(t *testing.T) {
  36. pid := int64(100)
  37. parent := &Account{ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
  38. shadow := &Account{ID: 200, ParentAccountID: &pid, QuotaDimension: QuotaDimensionSpark,
  39. Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Credentials: sparkCreds}
  40. s := newScheduler(map[int64]*Account{100: parent})
  41. require.True(t, s.isAccountRequestCompatible(ctx, shadow, sparkReq), "影子配 spark + 健康母 → 接 spark")
  42. require.False(t, s.isAccountRequestCompatible(ctx, shadow, normalReq), "影子(仅 spark mapping)→ 拒非 spark")
  43. })
  44. t.Run("empty_model_shadow_is_eligible_under_a2", func(t *testing.T) {
  45. // 有意的纯 A2 行为(用户裁决 2026-06-30):空 model 请求不经模型门过滤
  46. // (isAccountRequestCompatible 的 `req.RequestedModel != ""` 短路),故影子与普通账号
  47. // 一样成为候选。旧类型门曾在空 model 时排除影子(opt-in),该 opt-in 已随类型门移除——
  48. // routing 路径不再有任何类型判断。此测试锁定该决策,防被未来改动静默改回。
  49. pid := int64(100)
  50. parent := &Account{ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
  51. shadow := &Account{ID: 200, ParentAccountID: &pid, QuotaDimension: QuotaDimensionSpark,
  52. Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Credentials: sparkCreds}
  53. emptyReq := OpenAIAccountScheduleRequest{RequestedModel: "", Platform: PlatformOpenAI}
  54. s := newScheduler(map[int64]*Account{100: parent})
  55. require.True(t, s.isAccountRequestCompatible(ctx, shadow, emptyReq),
  56. "空 model 时影子可被选中(有意的纯 A2 行为:类型门移除后无 opt-in 排除)")
  57. })
  58. }
  59. // TestParentHealthSchedulerIntegration 通过 isAccountRequestCompatible 验证「母账号不可调度时影子被
  60. // 调度器拒绝」这一联动在调度器层面端到端生效。
  61. //
  62. // 使用的接缝:defaultOpenAIAccountScheduler.isAccountRequestCompatible,它通过
  63. // s.service.schedulerSnapshot.GetAccount(ctx, parentID) 解析母账号;
  64. // openAISnapshotCacheStub.accountsByID 提供对应的测试桩。
  65. func TestParentHealthSchedulerIntegration(t *testing.T) {
  66. ctx := context.Background()
  67. pid := int64(78100)
  68. sparkModel := "gpt-5.3-codex-spark"
  69. shadow := &Account{
  70. ID: 78200,
  71. ParentAccountID: &pid,
  72. QuotaDimension: QuotaDimensionSpark,
  73. Platform: PlatformOpenAI,
  74. Type: AccountTypeOAuth,
  75. Status: StatusActive,
  76. Schedulable: true,
  77. Concurrency: 1,
  78. }
  79. req := OpenAIAccountScheduleRequest{
  80. RequestedModel: sparkModel,
  81. Platform: PlatformOpenAI,
  82. }
  83. makeScheduler := func(parent *Account) *defaultOpenAIAccountScheduler {
  84. snapshotCache := &openAISnapshotCacheStub{
  85. accountsByID: map[int64]*Account{parent.ID: parent},
  86. }
  87. snapshotSvc := &SchedulerSnapshotService{cache: snapshotCache}
  88. svc := &OpenAIGatewayService{
  89. schedulerSnapshot: snapshotSvc,
  90. cfg: &config.Config{},
  91. }
  92. return &defaultOpenAIAccountScheduler{service: svc}
  93. }
  94. t.Run("unhealthy_parent_status_error_rejects_shadow", func(t *testing.T) {
  95. unhealthyParent := &Account{
  96. ID: 78100,
  97. Platform: PlatformOpenAI,
  98. Type: AccountTypeOAuth,
  99. Status: StatusError, // IsActive()==false → IsSchedulable()==false
  100. Schedulable: true,
  101. }
  102. require.False(t, unhealthyParent.IsSchedulable(), "前提:Status=error 的母账号不可调度")
  103. scheduler := makeScheduler(unhealthyParent)
  104. require.False(t, scheduler.isAccountRequestCompatible(ctx, shadow, req),
  105. "母账号不可调度时,影子账号必须被调度器拒绝")
  106. })
  107. t.Run("manual_schedulable_false_parent_does_not_reject_shadow", func(t *testing.T) {
  108. // F1 决策 A:母账号手动暂停(Schedulable=false)不传播到影子 —— 凭据仍可用,影子应被接受。
  109. manualPausedParent := &Account{
  110. ID: 78100,
  111. Platform: PlatformOpenAI,
  112. Type: AccountTypeOAuth,
  113. Status: StatusActive,
  114. Schedulable: false, // 显式手动暂停
  115. }
  116. require.False(t, manualPausedParent.IsSchedulable(), "前提:手动暂停的母账号自身不可调度")
  117. scheduler := makeScheduler(manualPausedParent)
  118. require.True(t, scheduler.isAccountRequestCompatible(ctx, shadow, req),
  119. "母账号手动暂停不应连坐影子(凭据仍可用)")
  120. })
  121. t.Run("global_rate_limited_parent_does_not_reject_shadow", func(t *testing.T) {
  122. // F1 核心修复:母账号 global 429(RateLimitResetAt)不连坐 spark 影子。
  123. resetAt := time.Now().Add(1 * time.Hour)
  124. rateLimitedParent := &Account{
  125. ID: 78100,
  126. Platform: PlatformOpenAI,
  127. Type: AccountTypeOAuth,
  128. Status: StatusActive,
  129. Schedulable: true,
  130. RateLimitResetAt: &resetAt,
  131. }
  132. require.False(t, rateLimitedParent.IsSchedulable(), "前提:global 限流母账号自身不可调度")
  133. scheduler := makeScheduler(rateLimitedParent)
  134. require.True(t, scheduler.isAccountRequestCompatible(ctx, shadow, req),
  135. "母账号 global 限流不应连坐 spark 影子")
  136. })
  137. t.Run("healthy_parent_accepts_shadow_control", func(t *testing.T) {
  138. healthyParent := &Account{
  139. ID: 78100,
  140. Platform: PlatformOpenAI,
  141. Type: AccountTypeOAuth,
  142. Status: StatusActive,
  143. Schedulable: true,
  144. }
  145. require.True(t, healthyParent.IsSchedulable(), "前提:健康母账号必须可调度")
  146. scheduler := makeScheduler(healthyParent)
  147. require.True(t, scheduler.isAccountRequestCompatible(ctx, shadow, req),
  148. "健康母账号时,影子账号必须被调度器接受(对照组)")
  149. })
  150. }
  151. func TestParentHealthSchedulerFallsBackToRepoWhenSnapshotMissesParent(t *testing.T) {
  152. ctx := context.Background()
  153. parentID := int64(79100)
  154. parent := Account{
  155. ID: parentID,
  156. Platform: PlatformOpenAI,
  157. Type: AccountTypeOAuth,
  158. Status: StatusActive,
  159. Schedulable: true,
  160. }
  161. shadow := &Account{
  162. ID: 79200,
  163. ParentAccountID: &parentID,
  164. QuotaDimension: QuotaDimensionSpark,
  165. Platform: PlatformOpenAI,
  166. Type: AccountTypeOAuth,
  167. Status: StatusActive,
  168. Schedulable: true,
  169. Concurrency: 1,
  170. }
  171. repo := schedulerTestOpenAIAccountRepo{accounts: []Account{parent}}
  172. scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{
  173. accountRepo: repo,
  174. schedulerSnapshot: &SchedulerSnapshotService{
  175. cache: &openAISnapshotCacheStub{},
  176. accountRepo: repo,
  177. cfg: &config.Config{
  178. Gateway: config.GatewayConfig{
  179. Scheduling: config.GatewaySchedulingConfig{
  180. DbFallbackEnabled: false,
  181. },
  182. },
  183. },
  184. },
  185. cfg: &config.Config{},
  186. }}
  187. require.True(t, scheduler.isAccountRequestCompatible(ctx, shadow, OpenAIAccountScheduleRequest{
  188. RequestedModel: "gpt-5.3-codex-spark",
  189. Platform: PlatformOpenAI,
  190. }), "快照缺失母账号且调度快照 DB fallback 关闭时,应回退 repo 解析健康母账号")
  191. }