scheduler_snapshot_bulk_event_test.go 8.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "sync"
  6. "testing"
  7. "github.com/Wei-Shaw/sub2api/internal/config"
  8. "github.com/stretchr/testify/require"
  9. )
  10. type bulkEventAccountRepo struct {
  11. *batchAccountQueryRepo
  12. accounts []*Account
  13. }
  14. func newBulkEventAccountRepo(accounts ...*Account) *bulkEventAccountRepo {
  15. return &bulkEventAccountRepo{
  16. batchAccountQueryRepo: newBatchAccountQueryRepo(),
  17. accounts: accounts,
  18. }
  19. }
  20. func (r *bulkEventAccountRepo) GetByIDs(context.Context, []int64) ([]*Account, error) {
  21. return append([]*Account(nil), r.accounts...), nil
  22. }
  23. type bulkEventSnapshotCache struct {
  24. *batchSnapshotCache
  25. accountMu sync.Mutex
  26. setAccountIDs []int64
  27. deleteAccountIDs []int64
  28. }
  29. func newBulkEventSnapshotCache() *bulkEventSnapshotCache {
  30. return &bulkEventSnapshotCache{batchSnapshotCache: newBatchSnapshotCache()}
  31. }
  32. func (c *bulkEventSnapshotCache) SetAccount(_ context.Context, account *Account) error {
  33. c.accountMu.Lock()
  34. defer c.accountMu.Unlock()
  35. c.setAccountIDs = append(c.setAccountIDs, account.ID)
  36. return nil
  37. }
  38. func (c *bulkEventSnapshotCache) DeleteAccount(_ context.Context, accountID int64) error {
  39. c.accountMu.Lock()
  40. defer c.accountMu.Unlock()
  41. c.deleteAccountIDs = append(c.deleteAccountIDs, accountID)
  42. return nil
  43. }
  44. func (c *bulkEventSnapshotCache) accountWrites() (set []int64, deleted []int64) {
  45. c.accountMu.Lock()
  46. defer c.accountMu.Unlock()
  47. return append([]int64(nil), c.setAccountIDs...), append([]int64(nil), c.deleteAccountIDs...)
  48. }
  49. func (c *bulkEventSnapshotCache) capturedBuckets() []SchedulerBucket {
  50. c.mu.Lock()
  51. defer c.mu.Unlock()
  52. return append([]SchedulerBucket(nil), c.captures...)
  53. }
  54. func newBulkEventTestService(cache SchedulerCache, accounts AccountRepository) *SchedulerSnapshotService {
  55. return NewSchedulerSnapshotService(cache, nil, accounts, nil, &config.Config{RunMode: config.RunModeStandard})
  56. }
  57. func bulkEventPayload(accountIDs []int64, groupIDs []int64) map[string]any {
  58. accountValues := make([]any, 0, len(accountIDs))
  59. for _, id := range accountIDs {
  60. accountValues = append(accountValues, id)
  61. }
  62. groupValues := make([]any, 0, len(groupIDs))
  63. for _, id := range groupIDs {
  64. groupValues = append(groupValues, id)
  65. }
  66. return map[string]any{
  67. "account_ids": accountValues,
  68. "group_ids": groupValues,
  69. }
  70. }
  71. func schedulerBucketsForTest(groupIDs []int64, platforms ...string) []SchedulerBucket {
  72. buckets := make([]SchedulerBucket, 0, len(groupIDs)*len(platforms)*3)
  73. for _, platform := range platforms {
  74. for _, groupID := range groupIDs {
  75. buckets = append(buckets,
  76. SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeSingle},
  77. SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeForced},
  78. )
  79. if platform == PlatformAnthropic || platform == PlatformGemini {
  80. buckets = append(buckets, SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeMixed})
  81. }
  82. }
  83. }
  84. return buckets
  85. }
  86. func TestSchedulerBulkAccountEventScopesOpenAIRebuildToFreshPlatform(t *testing.T) {
  87. cache := newBulkEventSnapshotCache()
  88. repo := newBulkEventAccountRepo(&Account{ID: 1, Platform: PlatformOpenAI, GroupIDs: []int64{12}})
  89. svc := newBulkEventTestService(cache, repo)
  90. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{1}, []int64{11}), make(map[batchSeenKey]struct{}))
  91. require.NoError(t, err)
  92. require.ElementsMatch(t, schedulerBucketsForTest([]int64{11, 12}, PlatformOpenAI), cache.capturedBuckets())
  93. set, deleted := cache.accountWrites()
  94. require.Equal(t, []int64{1}, set)
  95. require.Empty(t, deleted)
  96. }
  97. func TestSchedulerBulkAccountEventScopesCNRebuildToFreshPlatform(t *testing.T) {
  98. for _, platform := range []string{PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax} {
  99. t.Run(platform, func(t *testing.T) {
  100. cache := newBulkEventSnapshotCache()
  101. repo := newBulkEventAccountRepo(&Account{ID: 1, Platform: platform, GroupIDs: []int64{12}})
  102. svc := newBulkEventTestService(cache, repo)
  103. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{1}, []int64{11}), make(map[batchSeenKey]struct{}))
  104. require.NoError(t, err)
  105. require.ElementsMatch(t, schedulerBucketsForTest([]int64{11, 12}, platform), cache.capturedBuckets())
  106. })
  107. }
  108. }
  109. func TestSchedulerBulkAccountEventRebuildsOpenAIUngroupedBucket(t *testing.T) {
  110. cache := newBulkEventSnapshotCache()
  111. repo := newBulkEventAccountRepo(&Account{ID: 6, Platform: PlatformOpenAI})
  112. svc := newBulkEventTestService(cache, repo)
  113. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{6}, nil), make(map[batchSeenKey]struct{}))
  114. require.NoError(t, err)
  115. require.ElementsMatch(t, schedulerBucketsForTest([]int64{0}, PlatformOpenAI), cache.capturedBuckets())
  116. }
  117. func TestSchedulerBulkAccountEventKeepsGroupedAndUngroupedBuckets(t *testing.T) {
  118. cache := newBulkEventSnapshotCache()
  119. repo := newBulkEventAccountRepo(
  120. &Account{ID: 7, Platform: PlatformOpenAI, GroupIDs: []int64{51}},
  121. &Account{ID: 8, Platform: PlatformOpenAI},
  122. )
  123. svc := newBulkEventTestService(cache, repo)
  124. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{7, 8}, nil), make(map[batchSeenKey]struct{}))
  125. require.NoError(t, err)
  126. require.ElementsMatch(t, schedulerBucketsForTest([]int64{0, 51}, PlatformOpenAI), cache.capturedBuckets())
  127. }
  128. func TestSchedulerBulkAccountEventDoesNotCrossCurrentGroupsBetweenPlatforms(t *testing.T) {
  129. cache := newBulkEventSnapshotCache()
  130. repo := newBulkEventAccountRepo(
  131. &Account{ID: 9, Platform: PlatformOpenAI, GroupIDs: []int64{61}},
  132. &Account{ID: 10, Platform: PlatformGrok, GroupIDs: []int64{62}},
  133. )
  134. svc := newBulkEventTestService(cache, repo)
  135. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{9, 10}, []int64{63}), make(map[batchSeenKey]struct{}))
  136. require.NoError(t, err)
  137. want := append(
  138. schedulerBucketsForTest([]int64{61, 63}, PlatformOpenAI),
  139. schedulerBucketsForTest([]int64{62, 63}, PlatformGrok)...,
  140. )
  141. require.ElementsMatch(t, want, cache.capturedBuckets())
  142. }
  143. func TestSchedulerBulkAccountEventUsesGroupZeroInSimpleMode(t *testing.T) {
  144. cache := newBulkEventSnapshotCache()
  145. repo := newBulkEventAccountRepo(&Account{ID: 11, Platform: PlatformOpenAI, GroupIDs: []int64{71}})
  146. svc := NewSchedulerSnapshotService(cache, nil, repo, nil, &config.Config{RunMode: config.RunModeSimple})
  147. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{11}, []int64{72}), make(map[batchSeenKey]struct{}))
  148. require.NoError(t, err)
  149. require.ElementsMatch(t, schedulerBucketsForTest([]int64{0}, PlatformOpenAI), cache.capturedBuckets())
  150. }
  151. func TestSchedulerBulkAccountEventConservativelyExpandsAntigravityPlatforms(t *testing.T) {
  152. cache := newBulkEventSnapshotCache()
  153. // fresh 值可能已经关闭 mixed_scheduling,兼容平台仍要重建以清理旧快照。
  154. repo := newBulkEventAccountRepo(&Account{ID: 2, Platform: PlatformAntigravity, GroupIDs: []int64{22}})
  155. svc := newBulkEventTestService(cache, repo)
  156. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{2}, []int64{21}), make(map[batchSeenKey]struct{}))
  157. require.NoError(t, err)
  158. require.ElementsMatch(t,
  159. schedulerBucketsForTest([]int64{21, 22}, PlatformAnthropic, PlatformGemini, PlatformAntigravity),
  160. cache.capturedBuckets(),
  161. )
  162. }
  163. func TestSchedulerBulkAccountEventMissingAccountFallsBackToAllPlatforms(t *testing.T) {
  164. cache := newBulkEventSnapshotCache()
  165. repo := newBulkEventAccountRepo(&Account{ID: 3, Platform: PlatformOpenAI, GroupIDs: []int64{32}})
  166. svc := newBulkEventTestService(cache, repo)
  167. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{3, 4}, []int64{31}), make(map[batchSeenKey]struct{}))
  168. require.NoError(t, err)
  169. platforms := schedulerSnapshotPlatforms()
  170. require.ElementsMatch(t, schedulerBucketsForTest([]int64{31, 32}, platforms[:]...), cache.capturedBuckets())
  171. set, deleted := cache.accountWrites()
  172. require.Equal(t, []int64{3}, set)
  173. require.Equal(t, []int64{4}, deleted)
  174. }
  175. func TestSchedulerBulkAccountEventUnknownPlatformFallsBackToAllPlatforms(t *testing.T) {
  176. cache := newBulkEventSnapshotCache()
  177. repo := newBulkEventAccountRepo(&Account{ID: 5, GroupIDs: []int64{42}})
  178. svc := newBulkEventTestService(cache, repo)
  179. err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{5}, []int64{41}), make(map[batchSeenKey]struct{}))
  180. require.NoError(t, err)
  181. platforms := schedulerSnapshotPlatforms()
  182. require.ElementsMatch(t, schedulerBucketsForTest([]int64{41, 42}, platforms[:]...), cache.capturedBuckets())
  183. }