token_cache_invalidator_test.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "time"
  8. "github.com/stretchr/testify/require"
  9. )
  10. type geminiTokenCacheStub struct {
  11. deletedKeys []string
  12. deleteErr error
  13. }
  14. func (s *geminiTokenCacheStub) GetAccessToken(ctx context.Context, cacheKey string) (string, error) {
  15. return "", nil
  16. }
  17. func (s *geminiTokenCacheStub) SetAccessToken(ctx context.Context, cacheKey string, token string, ttl time.Duration) error {
  18. return nil
  19. }
  20. func (s *geminiTokenCacheStub) DeleteAccessToken(ctx context.Context, cacheKey string) error {
  21. s.deletedKeys = append(s.deletedKeys, cacheKey)
  22. return s.deleteErr
  23. }
  24. func (s *geminiTokenCacheStub) AcquireRefreshLock(ctx context.Context, cacheKey string, ttl time.Duration) (bool, error) {
  25. return true, nil
  26. }
  27. func (s *geminiTokenCacheStub) ReleaseRefreshLock(ctx context.Context, cacheKey string) error {
  28. return nil
  29. }
  30. func TestCompositeTokenCacheInvalidator_Gemini(t *testing.T) {
  31. cache := &geminiTokenCacheStub{}
  32. invalidator := NewCompositeTokenCacheInvalidator(cache)
  33. account := &Account{
  34. ID: 10,
  35. Platform: PlatformGemini,
  36. Type: AccountTypeOAuth,
  37. Credentials: map[string]any{
  38. "project_id": "project-x",
  39. },
  40. }
  41. err := invalidator.InvalidateToken(context.Background(), account)
  42. require.NoError(t, err)
  43. // 新行为:同时删除基于 project_id 和 account_id 的缓存键
  44. // 这是为了处理:首次获取 token 时可能没有 project_id,之后自动检测到后会使用新 key
  45. require.Equal(t, []string{"gemini:project-x", "gemini:account:10"}, cache.deletedKeys)
  46. }
  47. func TestCompositeTokenCacheInvalidator_GeminiWithoutProjectID(t *testing.T) {
  48. cache := &geminiTokenCacheStub{}
  49. invalidator := NewCompositeTokenCacheInvalidator(cache)
  50. account := &Account{
  51. ID: 10,
  52. Platform: PlatformGemini,
  53. Type: AccountTypeOAuth,
  54. Credentials: map[string]any{
  55. "access_token": "gemini-token",
  56. },
  57. }
  58. err := invalidator.InvalidateToken(context.Background(), account)
  59. require.NoError(t, err)
  60. // 没有 project_id 时,两个 key 相同,去重后只删除一个
  61. require.Equal(t, []string{"gemini:account:10"}, cache.deletedKeys)
  62. }
  63. func TestCompositeTokenCacheInvalidator_Antigravity(t *testing.T) {
  64. cache := &geminiTokenCacheStub{}
  65. invalidator := NewCompositeTokenCacheInvalidator(cache)
  66. account := &Account{
  67. ID: 99,
  68. Platform: PlatformAntigravity,
  69. Type: AccountTypeOAuth,
  70. Credentials: map[string]any{
  71. "project_id": "ag-project",
  72. },
  73. }
  74. err := invalidator.InvalidateToken(context.Background(), account)
  75. require.NoError(t, err)
  76. // 新行为:同时删除基于 project_id 和 account_id 的缓存键
  77. require.Equal(t, []string{"ag:ag-project", "ag:account:99"}, cache.deletedKeys)
  78. }
  79. func TestCompositeTokenCacheInvalidator_AntigravityWithoutProjectID(t *testing.T) {
  80. cache := &geminiTokenCacheStub{}
  81. invalidator := NewCompositeTokenCacheInvalidator(cache)
  82. account := &Account{
  83. ID: 99,
  84. Platform: PlatformAntigravity,
  85. Type: AccountTypeOAuth,
  86. Credentials: map[string]any{
  87. "access_token": "ag-token",
  88. },
  89. }
  90. err := invalidator.InvalidateToken(context.Background(), account)
  91. require.NoError(t, err)
  92. // 没有 project_id 时,两个 key 相同,去重后只删除一个
  93. require.Equal(t, []string{"ag:account:99"}, cache.deletedKeys)
  94. }
  95. func TestCompositeTokenCacheInvalidator_OpenAI(t *testing.T) {
  96. cache := &geminiTokenCacheStub{}
  97. invalidator := NewCompositeTokenCacheInvalidator(cache)
  98. account := &Account{
  99. ID: 500,
  100. Platform: PlatformOpenAI,
  101. Type: AccountTypeOAuth,
  102. Credentials: map[string]any{
  103. "access_token": "openai-token",
  104. },
  105. }
  106. err := invalidator.InvalidateToken(context.Background(), account)
  107. require.NoError(t, err)
  108. require.Equal(t, []string{"openai:account:500"}, cache.deletedKeys)
  109. }
  110. func TestCompositeTokenCacheInvalidator_Claude(t *testing.T) {
  111. cache := &geminiTokenCacheStub{}
  112. invalidator := NewCompositeTokenCacheInvalidator(cache)
  113. account := &Account{
  114. ID: 600,
  115. Platform: PlatformAnthropic,
  116. Type: AccountTypeOAuth,
  117. Credentials: map[string]any{
  118. "access_token": "claude-token",
  119. },
  120. }
  121. err := invalidator.InvalidateToken(context.Background(), account)
  122. require.NoError(t, err)
  123. require.Equal(t, []string{"claude:account:600"}, cache.deletedKeys)
  124. }
  125. func TestCompositeTokenCacheInvalidator_SkipNonOAuth(t *testing.T) {
  126. cache := &geminiTokenCacheStub{}
  127. invalidator := NewCompositeTokenCacheInvalidator(cache)
  128. tests := []struct {
  129. name string
  130. account *Account
  131. }{
  132. {
  133. name: "gemini_api_key",
  134. account: &Account{
  135. ID: 1,
  136. Platform: PlatformGemini,
  137. Type: AccountTypeAPIKey,
  138. },
  139. },
  140. {
  141. name: "openai_api_key",
  142. account: &Account{
  143. ID: 2,
  144. Platform: PlatformOpenAI,
  145. Type: AccountTypeAPIKey,
  146. },
  147. },
  148. {
  149. name: "claude_api_key",
  150. account: &Account{
  151. ID: 3,
  152. Platform: PlatformAnthropic,
  153. Type: AccountTypeAPIKey,
  154. },
  155. },
  156. {
  157. name: "claude_setup_token",
  158. account: &Account{
  159. ID: 4,
  160. Platform: PlatformAnthropic,
  161. Type: AccountTypeSetupToken,
  162. },
  163. },
  164. }
  165. for _, tt := range tests {
  166. t.Run(tt.name, func(t *testing.T) {
  167. cache.deletedKeys = nil
  168. err := invalidator.InvalidateToken(context.Background(), tt.account)
  169. require.NoError(t, err)
  170. require.Empty(t, cache.deletedKeys)
  171. })
  172. }
  173. }
  174. func TestCompositeTokenCacheInvalidator_SkipUnsupportedPlatform(t *testing.T) {
  175. cache := &geminiTokenCacheStub{}
  176. invalidator := NewCompositeTokenCacheInvalidator(cache)
  177. account := &Account{
  178. ID: 100,
  179. Platform: "unknown-platform",
  180. Type: AccountTypeOAuth,
  181. }
  182. err := invalidator.InvalidateToken(context.Background(), account)
  183. require.NoError(t, err)
  184. require.Empty(t, cache.deletedKeys)
  185. }
  186. func TestCompositeTokenCacheInvalidator_NilCache(t *testing.T) {
  187. invalidator := NewCompositeTokenCacheInvalidator(nil)
  188. account := &Account{
  189. ID: 2,
  190. Platform: PlatformGemini,
  191. Type: AccountTypeOAuth,
  192. }
  193. err := invalidator.InvalidateToken(context.Background(), account)
  194. require.NoError(t, err)
  195. }
  196. func TestCompositeTokenCacheInvalidator_NilAccount(t *testing.T) {
  197. cache := &geminiTokenCacheStub{}
  198. invalidator := NewCompositeTokenCacheInvalidator(cache)
  199. err := invalidator.InvalidateToken(context.Background(), nil)
  200. require.NoError(t, err)
  201. require.Empty(t, cache.deletedKeys)
  202. }
  203. func TestCompositeTokenCacheInvalidator_NilInvalidator(t *testing.T) {
  204. var invalidator *CompositeTokenCacheInvalidator
  205. account := &Account{
  206. ID: 5,
  207. Platform: PlatformGemini,
  208. Type: AccountTypeOAuth,
  209. }
  210. err := invalidator.InvalidateToken(context.Background(), account)
  211. require.NoError(t, err)
  212. }
  213. func TestCompositeTokenCacheInvalidator_DeleteError(t *testing.T) {
  214. expectedErr := errors.New("redis connection failed")
  215. cache := &geminiTokenCacheStub{deleteErr: expectedErr}
  216. invalidator := NewCompositeTokenCacheInvalidator(cache)
  217. tests := []struct {
  218. name string
  219. account *Account
  220. }{
  221. {
  222. name: "openai_delete_error",
  223. account: &Account{
  224. ID: 700,
  225. Platform: PlatformOpenAI,
  226. Type: AccountTypeOAuth,
  227. },
  228. },
  229. {
  230. name: "claude_delete_error",
  231. account: &Account{
  232. ID: 800,
  233. Platform: PlatformAnthropic,
  234. Type: AccountTypeOAuth,
  235. },
  236. },
  237. }
  238. for _, tt := range tests {
  239. t.Run(tt.name, func(t *testing.T) {
  240. // 新行为:删除失败只记录日志,不返回错误
  241. // 这是因为缓存失效失败不应影响主业务流程
  242. err := invalidator.InvalidateToken(context.Background(), tt.account)
  243. require.NoError(t, err)
  244. })
  245. }
  246. }
  247. func TestCompositeTokenCacheInvalidator_AllPlatformsIntegration(t *testing.T) {
  248. // 测试所有平台的缓存键生成和删除
  249. cache := &geminiTokenCacheStub{}
  250. invalidator := NewCompositeTokenCacheInvalidator(cache)
  251. accounts := []*Account{
  252. {ID: 1, Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "gemini-proj"}},
  253. {ID: 2, Platform: PlatformAntigravity, Type: AccountTypeOAuth, Credentials: map[string]any{"project_id": "ag-proj"}},
  254. {ID: 3, Platform: PlatformOpenAI, Type: AccountTypeOAuth},
  255. {ID: 4, Platform: PlatformAnthropic, Type: AccountTypeOAuth},
  256. }
  257. // 新行为:Gemini 和 Antigravity 会同时删除基于 project_id 和 account_id 的键
  258. expectedKeys := []string{
  259. "gemini:gemini-proj",
  260. "gemini:account:1",
  261. "ag:ag-proj",
  262. "ag:account:2",
  263. "openai:account:3",
  264. "claude:account:4",
  265. }
  266. for _, acc := range accounts {
  267. err := invalidator.InvalidateToken(context.Background(), acc)
  268. require.NoError(t, err)
  269. }
  270. require.Equal(t, expectedKeys, cache.deletedKeys)
  271. }
  272. // ========== GetCredentialAsInt64 测试 ==========
  273. func TestAccount_GetCredentialAsInt64(t *testing.T) {
  274. tests := []struct {
  275. name string
  276. credentials map[string]any
  277. key string
  278. expected int64
  279. }{
  280. {
  281. name: "int64_value",
  282. credentials: map[string]any{"_token_version": int64(1737654321000)},
  283. key: "_token_version",
  284. expected: 1737654321000,
  285. },
  286. {
  287. name: "float64_value",
  288. credentials: map[string]any{"_token_version": float64(1737654321000)},
  289. key: "_token_version",
  290. expected: 1737654321000,
  291. },
  292. {
  293. name: "int_value",
  294. credentials: map[string]any{"_token_version": 12345},
  295. key: "_token_version",
  296. expected: 12345,
  297. },
  298. {
  299. name: "string_value",
  300. credentials: map[string]any{"_token_version": "1737654321000"},
  301. key: "_token_version",
  302. expected: 1737654321000,
  303. },
  304. {
  305. name: "string_with_spaces",
  306. credentials: map[string]any{"_token_version": " 1737654321000 "},
  307. key: "_token_version",
  308. expected: 1737654321000,
  309. },
  310. {
  311. name: "nil_credentials",
  312. credentials: nil,
  313. key: "_token_version",
  314. expected: 0,
  315. },
  316. {
  317. name: "missing_key",
  318. credentials: map[string]any{"other_key": 123},
  319. key: "_token_version",
  320. expected: 0,
  321. },
  322. {
  323. name: "nil_value",
  324. credentials: map[string]any{"_token_version": nil},
  325. key: "_token_version",
  326. expected: 0,
  327. },
  328. {
  329. name: "invalid_string",
  330. credentials: map[string]any{"_token_version": "not_a_number"},
  331. key: "_token_version",
  332. expected: 0,
  333. },
  334. {
  335. name: "empty_string",
  336. credentials: map[string]any{"_token_version": ""},
  337. key: "_token_version",
  338. expected: 0,
  339. },
  340. }
  341. for _, tt := range tests {
  342. t.Run(tt.name, func(t *testing.T) {
  343. account := &Account{Credentials: tt.credentials}
  344. result := account.GetCredentialAsInt64(tt.key)
  345. require.Equal(t, tt.expected, result)
  346. })
  347. }
  348. }
  349. func TestAccount_GetCredentialAsInt64_NilAccount(t *testing.T) {
  350. var account *Account
  351. result := account.GetCredentialAsInt64("_token_version")
  352. require.Equal(t, int64(0), result)
  353. }
  354. // ========== CheckTokenVersion 测试 ==========
  355. func TestCheckTokenVersion(t *testing.T) {
  356. tests := []struct {
  357. name string
  358. account *Account
  359. latestAccount *Account
  360. repoErr error
  361. expectedStale bool
  362. }{
  363. {
  364. name: "nil_account",
  365. account: nil,
  366. latestAccount: nil,
  367. expectedStale: false,
  368. },
  369. {
  370. name: "no_version_in_account_but_db_has_version",
  371. account: &Account{
  372. ID: 1,
  373. Credentials: map[string]any{},
  374. },
  375. latestAccount: &Account{
  376. ID: 1,
  377. Credentials: map[string]any{"_token_version": int64(100)},
  378. },
  379. expectedStale: true, // 当前 account 无版本但 DB 有,说明已被异步刷新,当前已过时
  380. },
  381. {
  382. name: "both_no_version",
  383. account: &Account{
  384. ID: 1,
  385. Credentials: map[string]any{},
  386. },
  387. latestAccount: &Account{
  388. ID: 1,
  389. Credentials: map[string]any{},
  390. },
  391. expectedStale: false, // 两边都没有版本号,说明从未被异步刷新过,允许缓存
  392. },
  393. {
  394. name: "same_version",
  395. account: &Account{
  396. ID: 1,
  397. Credentials: map[string]any{"_token_version": int64(100)},
  398. },
  399. latestAccount: &Account{
  400. ID: 1,
  401. Credentials: map[string]any{"_token_version": int64(100)},
  402. },
  403. expectedStale: false,
  404. },
  405. {
  406. name: "current_version_newer",
  407. account: &Account{
  408. ID: 1,
  409. Credentials: map[string]any{"_token_version": int64(200)},
  410. },
  411. latestAccount: &Account{
  412. ID: 1,
  413. Credentials: map[string]any{"_token_version": int64(100)},
  414. },
  415. expectedStale: false,
  416. },
  417. {
  418. name: "current_version_older_stale",
  419. account: &Account{
  420. ID: 1,
  421. Credentials: map[string]any{"_token_version": int64(100)},
  422. },
  423. latestAccount: &Account{
  424. ID: 1,
  425. Credentials: map[string]any{"_token_version": int64(200)},
  426. },
  427. expectedStale: true, // 当前版本过时
  428. },
  429. {
  430. name: "repo_error",
  431. account: &Account{
  432. ID: 1,
  433. Credentials: map[string]any{"_token_version": int64(100)},
  434. },
  435. latestAccount: nil,
  436. repoErr: errors.New("db error"),
  437. expectedStale: false, // 查询失败,默认允许缓存
  438. },
  439. {
  440. name: "repo_returns_nil",
  441. account: &Account{
  442. ID: 1,
  443. Credentials: map[string]any{"_token_version": int64(100)},
  444. },
  445. latestAccount: nil,
  446. repoErr: nil,
  447. expectedStale: false, // 查询返回 nil,默认允许缓存
  448. },
  449. }
  450. for _, tt := range tests {
  451. t.Run(tt.name, func(t *testing.T) {
  452. // 由于 CheckTokenVersion 接受 AccountRepository 接口,而创建完整的 mock 很繁琐
  453. // 这里我们直接测试函数的核心逻辑来验证行为
  454. if tt.name == "nil_account" {
  455. _, isStale := CheckTokenVersion(context.Background(), nil, nil)
  456. require.Equal(t, tt.expectedStale, isStale)
  457. return
  458. }
  459. // 模拟 CheckTokenVersion 的核心逻辑
  460. account := tt.account
  461. currentVersion := account.GetCredentialAsInt64("_token_version")
  462. // 模拟 repo 查询
  463. latestAccount := tt.latestAccount
  464. if tt.repoErr != nil || latestAccount == nil {
  465. require.Equal(t, tt.expectedStale, false)
  466. return
  467. }
  468. latestVersion := latestAccount.GetCredentialAsInt64("_token_version")
  469. // 情况1: 当前 account 没有版本号,但 DB 中已有版本号
  470. if currentVersion == 0 && latestVersion > 0 {
  471. require.Equal(t, tt.expectedStale, true)
  472. return
  473. }
  474. // 情况2: 两边都没有版本号
  475. if currentVersion == 0 && latestVersion == 0 {
  476. require.Equal(t, tt.expectedStale, false)
  477. return
  478. }
  479. // 情况3: 比较版本号
  480. isStale := latestVersion > currentVersion
  481. require.Equal(t, tt.expectedStale, isStale)
  482. })
  483. }
  484. }
  485. func TestCheckTokenVersion_NilRepo(t *testing.T) {
  486. account := &Account{
  487. ID: 1,
  488. Credentials: map[string]any{"_token_version": int64(100)},
  489. }
  490. _, isStale := CheckTokenVersion(context.Background(), account, nil)
  491. require.False(t, isStale) // nil repo,默认允许缓存
  492. }