gemini_multiplatform_test.go 34 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  9. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  10. "github.com/stretchr/testify/require"
  11. )
  12. // mockAccountRepoForGemini Gemini 测试用的 mock
  13. type mockAccountRepoForGemini struct {
  14. accounts []Account
  15. accountsByID map[int64]*Account
  16. listByGroupFunc func(ctx context.Context, groupID int64, platforms []string) ([]Account, error)
  17. listByPlatformFunc func(ctx context.Context, platforms []string) ([]Account, error)
  18. }
  19. func (m *mockAccountRepoForGemini) GetByID(ctx context.Context, id int64) (*Account, error) {
  20. if acc, ok := m.accountsByID[id]; ok {
  21. return acc, nil
  22. }
  23. return nil, errors.New("account not found")
  24. }
  25. func (m *mockAccountRepoForGemini) GetByIDs(ctx context.Context, ids []int64) ([]*Account, error) {
  26. var result []*Account
  27. for _, id := range ids {
  28. if acc, ok := m.accountsByID[id]; ok {
  29. result = append(result, acc)
  30. }
  31. }
  32. return result, nil
  33. }
  34. func (m *mockAccountRepoForGemini) ExistsByID(ctx context.Context, id int64) (bool, error) {
  35. if m.accountsByID == nil {
  36. return false, nil
  37. }
  38. _, ok := m.accountsByID[id]
  39. return ok, nil
  40. }
  41. func (m *mockAccountRepoForGemini) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
  42. var result []Account
  43. for _, acc := range m.accounts {
  44. if acc.Platform == platform && acc.IsSchedulable() {
  45. result = append(result, acc)
  46. }
  47. }
  48. return result, nil
  49. }
  50. func (m *mockAccountRepoForGemini) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
  51. // 测试时不区分 groupID,直接按 platform 过滤
  52. return m.ListSchedulableByPlatform(ctx, platform)
  53. }
  54. // Stub methods to implement AccountRepository interface
  55. func (m *mockAccountRepoForGemini) Create(ctx context.Context, account *Account) error { return nil }
  56. func (m *mockAccountRepoForGemini) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
  57. return nil, nil
  58. }
  59. func (m *mockAccountRepoForGemini) FindByExtraField(ctx context.Context, key string, value any) ([]Account, error) {
  60. return nil, nil
  61. }
  62. func (m *mockAccountRepoForGemini) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
  63. return nil, nil
  64. }
  65. func (m *mockAccountRepoForGemini) Update(ctx context.Context, account *Account) error { return nil }
  66. func (m *mockAccountRepoForGemini) Delete(ctx context.Context, id int64) error { return nil }
  67. func (m *mockAccountRepoForGemini) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
  68. return nil, nil, nil
  69. }
  70. func (m *mockAccountRepoForGemini) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
  71. return nil, nil, nil
  72. }
  73. func (m *mockAccountRepoForGemini) ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error) {
  74. return nil, nil
  75. }
  76. func (m *mockAccountRepoForGemini) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) {
  77. return nil, nil
  78. }
  79. func (m *mockAccountRepoForGemini) ListActive(ctx context.Context) ([]Account, error) {
  80. return nil, nil
  81. }
  82. func (m *mockAccountRepoForGemini) ListByPlatform(ctx context.Context, platform string) ([]Account, error) {
  83. return nil, nil
  84. }
  85. func (m *mockAccountRepoForGemini) UpdateLastUsed(ctx context.Context, id int64) error { return nil }
  86. func (m *mockAccountRepoForGemini) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
  87. return nil
  88. }
  89. func (m *mockAccountRepoForGemini) SetError(ctx context.Context, id int64, errorMsg string) error {
  90. return nil
  91. }
  92. func (m *mockAccountRepoForGemini) ClearError(ctx context.Context, id int64) error {
  93. return nil
  94. }
  95. func (m *mockAccountRepoForGemini) SetSchedulable(ctx context.Context, id int64, schedulable bool) error {
  96. return nil
  97. }
  98. func (m *mockAccountRepoForGemini) AutoPauseExpiredAccounts(ctx context.Context, now time.Time) (int64, error) {
  99. return 0, nil
  100. }
  101. func (m *mockAccountRepoForGemini) BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error {
  102. return nil
  103. }
  104. func (m *mockAccountRepoForGemini) ListSchedulable(ctx context.Context) ([]Account, error) {
  105. return nil, nil
  106. }
  107. func (m *mockAccountRepoForGemini) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
  108. return nil, nil
  109. }
  110. func (m *mockAccountRepoForGemini) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]Account, error) {
  111. if m.listByPlatformFunc != nil {
  112. return m.listByPlatformFunc(ctx, platforms)
  113. }
  114. var result []Account
  115. platformSet := make(map[string]bool)
  116. for _, p := range platforms {
  117. platformSet[p] = true
  118. }
  119. for _, acc := range m.accounts {
  120. if platformSet[acc.Platform] && acc.IsSchedulable() {
  121. result = append(result, acc)
  122. }
  123. }
  124. return result, nil
  125. }
  126. func (m *mockAccountRepoForGemini) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error) {
  127. if m.listByGroupFunc != nil {
  128. return m.listByGroupFunc(ctx, groupID, platforms)
  129. }
  130. return m.ListSchedulableByPlatforms(ctx, platforms)
  131. }
  132. func (m *mockAccountRepoForGemini) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
  133. return m.ListSchedulableByPlatform(ctx, platform)
  134. }
  135. func (m *mockAccountRepoForGemini) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) {
  136. return m.ListSchedulableByPlatforms(ctx, platforms)
  137. }
  138. func (m *mockAccountRepoForGemini) ListModelAvailabilityCandidates(ctx context.Context, _ *int64, platforms []string, _ bool) ([]Account, error) {
  139. return m.ListSchedulableByPlatforms(ctx, platforms)
  140. }
  141. func (m *mockAccountRepoForGemini) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
  142. return nil
  143. }
  144. func (m *mockAccountRepoForGemini) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
  145. return nil
  146. }
  147. func (m *mockAccountRepoForGemini) SetOverloaded(ctx context.Context, id int64, until time.Time) error {
  148. return nil
  149. }
  150. func (m *mockAccountRepoForGemini) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
  151. return nil
  152. }
  153. func (m *mockAccountRepoForGemini) ClearTempUnschedulable(ctx context.Context, id int64) error {
  154. return nil
  155. }
  156. func (m *mockAccountRepoForGemini) ClearRateLimit(ctx context.Context, id int64) error { return nil }
  157. func (m *mockAccountRepoForGemini) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
  158. return nil
  159. }
  160. func (m *mockAccountRepoForGemini) ClearModelRateLimits(ctx context.Context, id int64) error {
  161. return nil
  162. }
  163. func (m *mockAccountRepoForGemini) UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error {
  164. return nil
  165. }
  166. func (m *mockAccountRepoForGemini) UpdateSessionWindowEnd(ctx context.Context, id int64, end time.Time) error {
  167. return nil
  168. }
  169. func (m *mockAccountRepoForGemini) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
  170. return nil
  171. }
  172. func (m *mockAccountRepoForGemini) BulkUpdate(ctx context.Context, ids []int64, updates AccountBulkUpdate) (int64, error) {
  173. return 0, nil
  174. }
  175. func (m *mockAccountRepoForGemini) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) error {
  176. return nil
  177. }
  178. func (m *mockAccountRepoForGemini) ResetQuotaUsedAndClearRateLimitCooldown(ctx context.Context, id int64) error {
  179. return nil
  180. }
  181. func (m *mockAccountRepoForGemini) RevertProxyFallback(ctx context.Context, accountID int64) error {
  182. return nil
  183. }
  184. func (m *mockAccountRepoForGemini) ListShadowsByParent(ctx context.Context, parentID int64) ([]*Account, error) {
  185. return nil, nil
  186. }
  187. // Verify interface implementation
  188. var _ AccountRepository = (*mockAccountRepoForGemini)(nil)
  189. // mockGroupRepoForGemini Gemini 测试用的 group repo mock
  190. type mockGroupRepoForGemini struct {
  191. groups map[int64]*Group
  192. getByIDCalls int
  193. getByIDLiteCalls int
  194. }
  195. func (m *mockGroupRepoForGemini) GetByID(ctx context.Context, id int64) (*Group, error) {
  196. m.getByIDCalls++
  197. if g, ok := m.groups[id]; ok {
  198. return g, nil
  199. }
  200. return nil, errors.New("group not found")
  201. }
  202. func (m *mockGroupRepoForGemini) GetByIDLite(ctx context.Context, id int64) (*Group, error) {
  203. m.getByIDLiteCalls++
  204. if g, ok := m.groups[id]; ok {
  205. return g, nil
  206. }
  207. return nil, errors.New("group not found")
  208. }
  209. // Stub methods to implement GroupRepository interface
  210. func (m *mockGroupRepoForGemini) Create(ctx context.Context, group *Group) error { return nil }
  211. func (m *mockGroupRepoForGemini) Update(ctx context.Context, group *Group) error { return nil }
  212. func (m *mockGroupRepoForGemini) Delete(ctx context.Context, id int64) error { return nil }
  213. func (m *mockGroupRepoForGemini) DeleteCascade(ctx context.Context, id int64) ([]int64, error) {
  214. return nil, nil
  215. }
  216. func (m *mockGroupRepoForGemini) List(ctx context.Context, params pagination.PaginationParams) ([]Group, *pagination.PaginationResult, error) {
  217. return nil, nil, nil
  218. }
  219. func (m *mockGroupRepoForGemini) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]Group, *pagination.PaginationResult, error) {
  220. return nil, nil, nil
  221. }
  222. func (m *mockGroupRepoForGemini) ListActive(ctx context.Context) ([]Group, error) { return nil, nil }
  223. func (m *mockGroupRepoForGemini) ListActiveByPlatform(ctx context.Context, platform string) ([]Group, error) {
  224. return nil, nil
  225. }
  226. func (m *mockGroupRepoForGemini) ExistsByName(ctx context.Context, name string) (bool, error) {
  227. return false, nil
  228. }
  229. func (m *mockGroupRepoForGemini) GetAccountCount(ctx context.Context, groupID int64) (int64, int64, error) {
  230. return 0, 0, nil
  231. }
  232. func (m *mockGroupRepoForGemini) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
  233. return 0, nil
  234. }
  235. func (m *mockGroupRepoForGemini) BindAccountsToGroup(ctx context.Context, groupID int64, accountIDs []int64) error {
  236. return nil
  237. }
  238. func (m *mockGroupRepoForGemini) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int64) ([]int64, error) {
  239. return nil, nil
  240. }
  241. func (m *mockGroupRepoForGemini) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
  242. return nil
  243. }
  244. var _ GroupRepository = (*mockGroupRepoForGemini)(nil)
  245. // mockGatewayCacheForGemini Gemini 测试用的 cache mock
  246. type mockGatewayCacheForGemini struct {
  247. sessionBindings map[string]int64
  248. deletedSessions map[string]int
  249. }
  250. func (m *mockGatewayCacheForGemini) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
  251. if id, ok := m.sessionBindings[sessionHash]; ok {
  252. return id, nil
  253. }
  254. return 0, errors.New("not found")
  255. }
  256. func (m *mockGatewayCacheForGemini) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
  257. if m.sessionBindings == nil {
  258. m.sessionBindings = make(map[string]int64)
  259. }
  260. m.sessionBindings[sessionHash] = accountID
  261. return nil
  262. }
  263. func (m *mockGatewayCacheForGemini) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
  264. return nil
  265. }
  266. func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
  267. if m.sessionBindings == nil {
  268. return nil
  269. }
  270. if m.deletedSessions == nil {
  271. m.deletedSessions = make(map[string]int)
  272. }
  273. m.deletedSessions[sessionHash]++
  274. delete(m.sessionBindings, sessionHash)
  275. return nil
  276. }
  277. func (m *mockGatewayCacheForGemini) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
  278. return nil
  279. }
  280. func (m *mockGatewayCacheForGemini) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
  281. return nil, nil
  282. }
  283. func (m *mockGatewayCacheForGemini) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
  284. return true, nil
  285. }
  286. func (m *mockGatewayCacheForGemini) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
  287. return nil
  288. }
  289. func (m *mockGatewayCacheForGemini) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
  290. return nil
  291. }
  292. func (m *mockGatewayCacheForGemini) GetReasoningContent(_ context.Context, _ string) (string, error) {
  293. return "", ErrReasoningContentNotFound
  294. }
  295. // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择
  296. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) {
  297. ctx := context.Background()
  298. repo := &mockAccountRepoForGemini{
  299. accounts: []Account{
  300. {ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  301. {ID: 2, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
  302. {ID: 3, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被隔离
  303. },
  304. accountsByID: map[int64]*Account{},
  305. }
  306. for i := range repo.accounts {
  307. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  308. }
  309. cache := &mockGatewayCacheForGemini{}
  310. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  311. svc := &GeminiMessagesCompatService{
  312. accountRepo: repo,
  313. groupRepo: groupRepo,
  314. cache: cache,
  315. }
  316. // 无分组时使用 gemini 平台
  317. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
  318. require.NoError(t, err)
  319. require.NotNil(t, acc)
  320. require.Equal(t, int64(1), acc.ID, "应选择优先级最高的 gemini 账户")
  321. require.Equal(t, PlatformGemini, acc.Platform, "无分组时应只返回 gemini 平台账户")
  322. }
  323. func TestGeminiMessagesCompatService_GroupResolution_ReusesContextGroup(t *testing.T) {
  324. ctx := context.Background()
  325. groupID := int64(7)
  326. group := &Group{
  327. ID: groupID,
  328. Platform: PlatformGemini,
  329. Status: StatusActive,
  330. Hydrated: true,
  331. }
  332. ctx = context.WithValue(ctx, ctxkey.Group, group)
  333. repo := &mockAccountRepoForGemini{
  334. accounts: []Account{
  335. {ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  336. },
  337. accountsByID: map[int64]*Account{},
  338. }
  339. for i := range repo.accounts {
  340. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  341. }
  342. cache := &mockGatewayCacheForGemini{}
  343. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  344. svc := &GeminiMessagesCompatService{
  345. accountRepo: repo,
  346. groupRepo: groupRepo,
  347. cache: cache,
  348. }
  349. acc, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gemini-2.5-flash", nil)
  350. require.NoError(t, err)
  351. require.NotNil(t, acc)
  352. require.Equal(t, 0, groupRepo.getByIDCalls)
  353. require.Equal(t, 0, groupRepo.getByIDLiteCalls)
  354. }
  355. func TestGeminiMessagesCompatService_GroupResolution_UsesLiteFetch(t *testing.T) {
  356. ctx := context.Background()
  357. groupID := int64(7)
  358. repo := &mockAccountRepoForGemini{
  359. accounts: []Account{
  360. {ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  361. },
  362. accountsByID: map[int64]*Account{},
  363. }
  364. for i := range repo.accounts {
  365. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  366. }
  367. cache := &mockGatewayCacheForGemini{}
  368. groupRepo := &mockGroupRepoForGemini{
  369. groups: map[int64]*Group{
  370. groupID: {ID: groupID, Platform: PlatformGemini},
  371. },
  372. }
  373. svc := &GeminiMessagesCompatService{
  374. accountRepo: repo,
  375. groupRepo: groupRepo,
  376. cache: cache,
  377. }
  378. acc, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gemini-2.5-flash", nil)
  379. require.NoError(t, err)
  380. require.NotNil(t, acc)
  381. require.Equal(t, 0, groupRepo.getByIDCalls)
  382. require.Equal(t, 1, groupRepo.getByIDLiteCalls)
  383. }
  384. // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_AntigravityGroup 测试 antigravity 分组
  385. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_AntigravityGroup(t *testing.T) {
  386. ctx := context.Background()
  387. repo := &mockAccountRepoForGemini{
  388. accounts: []Account{
  389. {ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被隔离
  390. {ID: 2, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被选择
  391. },
  392. accountsByID: map[int64]*Account{},
  393. }
  394. for i := range repo.accounts {
  395. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  396. }
  397. cache := &mockGatewayCacheForGemini{}
  398. groupRepo := &mockGroupRepoForGemini{
  399. groups: map[int64]*Group{
  400. 1: {ID: 1, Platform: PlatformAntigravity},
  401. },
  402. }
  403. svc := &GeminiMessagesCompatService{
  404. accountRepo: repo,
  405. groupRepo: groupRepo,
  406. cache: cache,
  407. }
  408. groupID := int64(1)
  409. acc, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gemini-2.5-flash", nil)
  410. require.NoError(t, err)
  411. require.NotNil(t, acc)
  412. require.Equal(t, int64(2), acc.ID)
  413. require.Equal(t, PlatformAntigravity, acc.Platform, "antigravity 分组应只返回 antigravity 账户")
  414. }
  415. // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_OAuthPreferred 测试 OAuth 优先
  416. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_OAuthPreferred(t *testing.T) {
  417. ctx := context.Background()
  418. repo := &mockAccountRepoForGemini{
  419. accounts: []Account{
  420. {ID: 1, Platform: PlatformGemini, Type: AccountTypeAPIKey, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: nil},
  421. {ID: 2, Platform: PlatformGemini, Type: AccountTypeOAuth, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: nil},
  422. },
  423. accountsByID: map[int64]*Account{},
  424. }
  425. for i := range repo.accounts {
  426. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  427. }
  428. cache := &mockGatewayCacheForGemini{}
  429. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  430. svc := &GeminiMessagesCompatService{
  431. accountRepo: repo,
  432. groupRepo: groupRepo,
  433. cache: cache,
  434. }
  435. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
  436. require.NoError(t, err)
  437. require.NotNil(t, acc)
  438. require.Equal(t, int64(2), acc.ID, "同优先级且都未使用时,应优先选择 OAuth 账户")
  439. require.Equal(t, AccountTypeOAuth, acc.Type)
  440. }
  441. // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_NoAvailableAccounts 测试无可用账户
  442. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_NoAvailableAccounts(t *testing.T) {
  443. ctx := context.Background()
  444. repo := &mockAccountRepoForGemini{
  445. accounts: []Account{},
  446. accountsByID: map[int64]*Account{},
  447. }
  448. cache := &mockGatewayCacheForGemini{}
  449. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  450. svc := &GeminiMessagesCompatService{
  451. accountRepo: repo,
  452. groupRepo: groupRepo,
  453. cache: cache,
  454. }
  455. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
  456. require.Error(t, err)
  457. require.Nil(t, acc)
  458. require.Contains(t, err.Error(), "no available")
  459. }
  460. // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_StickySession 测试粘性会话
  461. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_StickySession(t *testing.T) {
  462. ctx := context.Background()
  463. t.Run("粘性会话命中-同平台", func(t *testing.T) {
  464. repo := &mockAccountRepoForGemini{
  465. accounts: []Account{
  466. {ID: 1, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
  467. {ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  468. },
  469. accountsByID: map[int64]*Account{},
  470. }
  471. for i := range repo.accounts {
  472. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  473. }
  474. // 注意:缓存键使用 "gemini:" 前缀
  475. cache := &mockGatewayCacheForGemini{
  476. sessionBindings: map[string]int64{"gemini:session-123": 1},
  477. }
  478. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  479. svc := &GeminiMessagesCompatService{
  480. accountRepo: repo,
  481. groupRepo: groupRepo,
  482. cache: cache,
  483. }
  484. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-123", "gemini-2.5-flash", nil)
  485. require.NoError(t, err)
  486. require.NotNil(t, acc)
  487. require.Equal(t, int64(1), acc.ID, "应返回粘性会话绑定的账户")
  488. })
  489. t.Run("粘性会话平台不匹配-降级选择", func(t *testing.T) {
  490. repo := &mockAccountRepoForGemini{
  491. accounts: []Account{
  492. {ID: 1, Platform: PlatformAntigravity, Priority: 2, Status: StatusActive, Schedulable: true}, // 粘性会话绑定
  493. {ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  494. },
  495. accountsByID: map[int64]*Account{},
  496. }
  497. for i := range repo.accounts {
  498. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  499. }
  500. cache := &mockGatewayCacheForGemini{
  501. sessionBindings: map[string]int64{"gemini:session-123": 1}, // 绑定 antigravity 账户
  502. }
  503. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  504. svc := &GeminiMessagesCompatService{
  505. accountRepo: repo,
  506. groupRepo: groupRepo,
  507. cache: cache,
  508. }
  509. // 无分组时使用 gemini 平台,粘性会话绑定的 antigravity 账户平台不匹配
  510. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-123", "gemini-2.5-flash", nil)
  511. require.NoError(t, err)
  512. require.NotNil(t, acc)
  513. require.Equal(t, int64(2), acc.ID, "粘性会话账户平台不匹配,应降级选择 gemini 账户")
  514. require.Equal(t, PlatformGemini, acc.Platform)
  515. })
  516. t.Run("粘性会话不命中无前缀缓存键", func(t *testing.T) {
  517. repo := &mockAccountRepoForGemini{
  518. accounts: []Account{
  519. {ID: 1, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
  520. {ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  521. },
  522. accountsByID: map[int64]*Account{},
  523. }
  524. for i := range repo.accounts {
  525. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  526. }
  527. // 缓存键没有 "gemini:" 前缀,不应命中
  528. cache := &mockGatewayCacheForGemini{
  529. sessionBindings: map[string]int64{"session-123": 1},
  530. }
  531. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  532. svc := &GeminiMessagesCompatService{
  533. accountRepo: repo,
  534. groupRepo: groupRepo,
  535. cache: cache,
  536. }
  537. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-123", "gemini-2.5-flash", nil)
  538. require.NoError(t, err)
  539. require.NotNil(t, acc)
  540. // 粘性会话未命中,按优先级选择
  541. require.Equal(t, int64(2), acc.ID, "粘性会话未命中,应按优先级选择")
  542. })
  543. t.Run("粘性会话不可调度-清理并回退选择", func(t *testing.T) {
  544. repo := &mockAccountRepoForGemini{
  545. accounts: []Account{
  546. {ID: 1, Platform: PlatformGemini, Priority: 2, Status: StatusDisabled, Schedulable: true},
  547. {ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  548. },
  549. accountsByID: map[int64]*Account{},
  550. }
  551. for i := range repo.accounts {
  552. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  553. }
  554. cache := &mockGatewayCacheForGemini{
  555. sessionBindings: map[string]int64{"gemini:session-123": 1},
  556. }
  557. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  558. svc := &GeminiMessagesCompatService{
  559. accountRepo: repo,
  560. groupRepo: groupRepo,
  561. cache: cache,
  562. }
  563. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-123", "gemini-2.5-flash", nil)
  564. require.NoError(t, err)
  565. require.NotNil(t, acc)
  566. require.Equal(t, int64(2), acc.ID)
  567. require.Equal(t, 1, cache.deletedSessions["gemini:session-123"])
  568. require.Equal(t, int64(2), cache.sessionBindings["gemini:session-123"])
  569. })
  570. }
  571. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_ForcePlatformFallback(t *testing.T) {
  572. ctx := context.Background()
  573. groupID := int64(9)
  574. ctx = context.WithValue(ctx, ctxkey.ForcePlatform, PlatformAntigravity)
  575. repo := &mockAccountRepoForGemini{
  576. listByGroupFunc: func(ctx context.Context, groupID int64, platforms []string) ([]Account, error) {
  577. return nil, nil
  578. },
  579. listByPlatformFunc: func(ctx context.Context, platforms []string) ([]Account, error) {
  580. return []Account{
  581. {ID: 1, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true},
  582. }, nil
  583. },
  584. accountsByID: map[int64]*Account{
  585. 1: {ID: 1, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true},
  586. },
  587. }
  588. cache := &mockGatewayCacheForGemini{}
  589. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  590. svc := &GeminiMessagesCompatService{
  591. accountRepo: repo,
  592. groupRepo: groupRepo,
  593. cache: cache,
  594. }
  595. acc, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gemini-2.5-flash", nil)
  596. require.NoError(t, err)
  597. require.NotNil(t, acc)
  598. require.Equal(t, int64(1), acc.ID)
  599. }
  600. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_NoModelSupport(t *testing.T) {
  601. ctx := context.Background()
  602. repo := &mockAccountRepoForGemini{
  603. accounts: []Account{
  604. {
  605. ID: 1,
  606. Platform: PlatformGemini,
  607. Priority: 1,
  608. Status: StatusActive,
  609. Schedulable: true,
  610. Credentials: map[string]any{"model_mapping": map[string]any{"gemini-1.0-pro": "gemini-1.0-pro"}},
  611. },
  612. },
  613. accountsByID: map[int64]*Account{},
  614. }
  615. for i := range repo.accounts {
  616. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  617. }
  618. cache := &mockGatewayCacheForGemini{}
  619. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  620. svc := &GeminiMessagesCompatService{
  621. accountRepo: repo,
  622. groupRepo: groupRepo,
  623. cache: cache,
  624. }
  625. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
  626. require.Error(t, err)
  627. require.Nil(t, acc)
  628. require.Contains(t, err.Error(), "supporting model")
  629. }
  630. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_StickyMixedScheduling(t *testing.T) {
  631. ctx := context.Background()
  632. repo := &mockAccountRepoForGemini{
  633. accounts: []Account{
  634. {ID: 1, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true, Extra: map[string]any{"mixed_scheduling": true}},
  635. {ID: 2, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
  636. },
  637. accountsByID: map[int64]*Account{},
  638. }
  639. for i := range repo.accounts {
  640. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  641. }
  642. cache := &mockGatewayCacheForGemini{
  643. sessionBindings: map[string]int64{"gemini:session-999": 1},
  644. }
  645. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  646. svc := &GeminiMessagesCompatService{
  647. accountRepo: repo,
  648. groupRepo: groupRepo,
  649. cache: cache,
  650. }
  651. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-999", "gemini-2.5-flash", nil)
  652. require.NoError(t, err)
  653. require.NotNil(t, acc)
  654. require.Equal(t, int64(1), acc.ID)
  655. }
  656. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_SkipDisabledMixedScheduling(t *testing.T) {
  657. ctx := context.Background()
  658. repo := &mockAccountRepoForGemini{
  659. accounts: []Account{
  660. {ID: 1, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true},
  661. {ID: 2, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
  662. },
  663. accountsByID: map[int64]*Account{},
  664. }
  665. for i := range repo.accounts {
  666. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  667. }
  668. cache := &mockGatewayCacheForGemini{}
  669. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  670. svc := &GeminiMessagesCompatService{
  671. accountRepo: repo,
  672. groupRepo: groupRepo,
  673. cache: cache,
  674. }
  675. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
  676. require.NoError(t, err)
  677. require.NotNil(t, acc)
  678. require.Equal(t, int64(2), acc.ID)
  679. }
  680. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_ExcludedAccount(t *testing.T) {
  681. ctx := context.Background()
  682. repo := &mockAccountRepoForGemini{
  683. accounts: []Account{
  684. {ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
  685. {ID: 2, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
  686. },
  687. accountsByID: map[int64]*Account{},
  688. }
  689. for i := range repo.accounts {
  690. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  691. }
  692. cache := &mockGatewayCacheForGemini{}
  693. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  694. svc := &GeminiMessagesCompatService{
  695. accountRepo: repo,
  696. groupRepo: groupRepo,
  697. cache: cache,
  698. }
  699. excluded := map[int64]struct{}{1: {}}
  700. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", excluded)
  701. require.NoError(t, err)
  702. require.NotNil(t, acc)
  703. require.Equal(t, int64(2), acc.ID)
  704. }
  705. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_ListError(t *testing.T) {
  706. ctx := context.Background()
  707. repo := &mockAccountRepoForGemini{
  708. listByPlatformFunc: func(ctx context.Context, platforms []string) ([]Account, error) {
  709. return nil, errors.New("query failed")
  710. },
  711. }
  712. cache := &mockGatewayCacheForGemini{}
  713. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  714. svc := &GeminiMessagesCompatService{
  715. accountRepo: repo,
  716. groupRepo: groupRepo,
  717. cache: cache,
  718. }
  719. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
  720. require.Error(t, err)
  721. require.Nil(t, acc)
  722. require.Contains(t, err.Error(), "query accounts failed")
  723. }
  724. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_PreferOAuth(t *testing.T) {
  725. ctx := context.Background()
  726. repo := &mockAccountRepoForGemini{
  727. accounts: []Account{
  728. {ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true, Type: AccountTypeAPIKey},
  729. {ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true, Type: AccountTypeOAuth},
  730. },
  731. accountsByID: map[int64]*Account{},
  732. }
  733. for i := range repo.accounts {
  734. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  735. }
  736. cache := &mockGatewayCacheForGemini{}
  737. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  738. svc := &GeminiMessagesCompatService{
  739. accountRepo: repo,
  740. groupRepo: groupRepo,
  741. cache: cache,
  742. }
  743. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-pro", nil)
  744. require.NoError(t, err)
  745. require.NotNil(t, acc)
  746. require.Equal(t, int64(2), acc.ID)
  747. }
  748. func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_PreferLeastRecentlyUsed(t *testing.T) {
  749. ctx := context.Background()
  750. oldTime := time.Now().Add(-2 * time.Hour)
  751. newTime := time.Now().Add(-1 * time.Hour)
  752. repo := &mockAccountRepoForGemini{
  753. accounts: []Account{
  754. {ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: &newTime},
  755. {ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: &oldTime},
  756. },
  757. accountsByID: map[int64]*Account{},
  758. }
  759. for i := range repo.accounts {
  760. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  761. }
  762. cache := &mockGatewayCacheForGemini{}
  763. groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
  764. svc := &GeminiMessagesCompatService{
  765. accountRepo: repo,
  766. groupRepo: groupRepo,
  767. cache: cache,
  768. }
  769. acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-pro", nil)
  770. require.NoError(t, err)
  771. require.NotNil(t, acc)
  772. require.Equal(t, int64(2), acc.ID)
  773. }
  774. // TestGeminiPlatformRouting_DocumentRouteDecision 测试平台路由决策逻辑
  775. func TestGeminiPlatformRouting_DocumentRouteDecision(t *testing.T) {
  776. tests := []struct {
  777. name string
  778. platform string
  779. expectedService string // "gemini" 表示 ForwardNative, "antigravity" 表示 ForwardGemini
  780. }{
  781. {
  782. name: "Gemini平台走ForwardNative",
  783. platform: PlatformGemini,
  784. expectedService: "gemini",
  785. },
  786. {
  787. name: "Antigravity平台走ForwardGemini",
  788. platform: PlatformAntigravity,
  789. expectedService: "antigravity",
  790. },
  791. }
  792. for _, tt := range tests {
  793. t.Run(tt.name, func(t *testing.T) {
  794. account := &Account{Platform: tt.platform}
  795. // 模拟 Handler 层的路由逻辑
  796. var serviceName string
  797. if account.Platform == PlatformAntigravity {
  798. serviceName = "antigravity"
  799. } else {
  800. serviceName = "gemini"
  801. }
  802. require.Equal(t, tt.expectedService, serviceName,
  803. "平台 %s 应该路由到 %s 服务", tt.platform, tt.expectedService)
  804. })
  805. }
  806. }
  807. func TestGeminiMessagesCompatService_isModelSupportedByAccount(t *testing.T) {
  808. svc := &GeminiMessagesCompatService{}
  809. tests := []struct {
  810. name string
  811. account *Account
  812. model string
  813. expected bool
  814. }{
  815. {
  816. name: "Antigravity平台-支持gemini模型",
  817. account: &Account{Platform: PlatformAntigravity},
  818. model: "gemini-2.5-flash",
  819. expected: true,
  820. },
  821. {
  822. name: "Antigravity平台-支持claude模型",
  823. account: &Account{Platform: PlatformAntigravity},
  824. model: "claude-sonnet-4-5",
  825. expected: true,
  826. },
  827. {
  828. name: "Antigravity平台-不支持gpt模型",
  829. account: &Account{Platform: PlatformAntigravity},
  830. model: "gpt-4",
  831. expected: false,
  832. },
  833. {
  834. name: "Antigravity平台-空模型允许",
  835. account: &Account{Platform: PlatformAntigravity},
  836. model: "",
  837. expected: true,
  838. },
  839. {
  840. name: "Antigravity平台-自定义映射-支持自定义模型",
  841. account: &Account{
  842. Platform: PlatformAntigravity,
  843. Credentials: map[string]any{
  844. "model_mapping": map[string]any{
  845. "my-custom-model": "upstream-model",
  846. "gpt-4o": "some-model",
  847. },
  848. },
  849. },
  850. model: "my-custom-model",
  851. expected: true,
  852. },
  853. {
  854. name: "Antigravity平台-自定义映射-不在映射中的模型不支持",
  855. account: &Account{
  856. Platform: PlatformAntigravity,
  857. Credentials: map[string]any{
  858. "model_mapping": map[string]any{
  859. "my-custom-model": "upstream-model",
  860. },
  861. },
  862. },
  863. model: "claude-sonnet-4-5",
  864. expected: false,
  865. },
  866. {
  867. name: "Gemini平台-无映射配置-支持所有模型",
  868. account: &Account{Platform: PlatformGemini},
  869. model: "gemini-2.5-flash",
  870. expected: true,
  871. },
  872. {
  873. name: "Gemini平台-有映射配置-只支持配置的模型",
  874. account: &Account{
  875. Platform: PlatformGemini,
  876. Credentials: map[string]any{"model_mapping": map[string]any{"gemini-2.5-pro": "x"}},
  877. },
  878. model: "gemini-2.5-flash",
  879. expected: false,
  880. },
  881. }
  882. for _, tt := range tests {
  883. t.Run(tt.name, func(t *testing.T) {
  884. got := svc.isModelSupportedByAccount(tt.account, tt.model)
  885. require.Equal(t, tt.expected, got)
  886. })
  887. }
  888. }