openai_token_provider_test.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "sync"
  7. "sync/atomic"
  8. "testing"
  9. "time"
  10. "github.com/stretchr/testify/require"
  11. )
  12. // openAITokenCacheStub implements OpenAITokenCache for testing
  13. type openAITokenCacheStub struct {
  14. mu sync.Mutex
  15. tokens map[string]string
  16. getErr error
  17. setErr error
  18. deleteErr error
  19. lockAcquired bool
  20. lockErr error
  21. releaseLockErr error
  22. getCalled int32
  23. setCalled int32
  24. lockCalled int32
  25. unlockCalled int32
  26. simulateLockRace bool
  27. }
  28. func newOpenAITokenCacheStub() *openAITokenCacheStub {
  29. return &openAITokenCacheStub{
  30. tokens: make(map[string]string),
  31. lockAcquired: true,
  32. }
  33. }
  34. func (s *openAITokenCacheStub) GetAccessToken(ctx context.Context, cacheKey string) (string, error) {
  35. atomic.AddInt32(&s.getCalled, 1)
  36. if s.getErr != nil {
  37. return "", s.getErr
  38. }
  39. s.mu.Lock()
  40. defer s.mu.Unlock()
  41. return s.tokens[cacheKey], nil
  42. }
  43. func (s *openAITokenCacheStub) SetAccessToken(ctx context.Context, cacheKey string, token string, ttl time.Duration) error {
  44. atomic.AddInt32(&s.setCalled, 1)
  45. if s.setErr != nil {
  46. return s.setErr
  47. }
  48. s.mu.Lock()
  49. defer s.mu.Unlock()
  50. s.tokens[cacheKey] = token
  51. return nil
  52. }
  53. func (s *openAITokenCacheStub) DeleteAccessToken(ctx context.Context, cacheKey string) error {
  54. if s.deleteErr != nil {
  55. return s.deleteErr
  56. }
  57. s.mu.Lock()
  58. defer s.mu.Unlock()
  59. delete(s.tokens, cacheKey)
  60. return nil
  61. }
  62. func (s *openAITokenCacheStub) AcquireRefreshLock(ctx context.Context, cacheKey string, ttl time.Duration) (bool, error) {
  63. atomic.AddInt32(&s.lockCalled, 1)
  64. if s.lockErr != nil {
  65. return false, s.lockErr
  66. }
  67. if s.simulateLockRace {
  68. return false, nil
  69. }
  70. return s.lockAcquired, nil
  71. }
  72. func (s *openAITokenCacheStub) ReleaseRefreshLock(ctx context.Context, cacheKey string) error {
  73. atomic.AddInt32(&s.unlockCalled, 1)
  74. return s.releaseLockErr
  75. }
  76. // openAIAccountRepoStub is a minimal stub implementing only the methods used by OpenAITokenProvider
  77. type openAIAccountRepoStub struct {
  78. account *Account
  79. getErr error
  80. updateErr error
  81. getCalled int32
  82. updateCalled int32
  83. }
  84. func (r *openAIAccountRepoStub) GetByID(ctx context.Context, id int64) (*Account, error) {
  85. atomic.AddInt32(&r.getCalled, 1)
  86. if r.getErr != nil {
  87. return nil, r.getErr
  88. }
  89. return r.account, nil
  90. }
  91. func (r *openAIAccountRepoStub) Update(ctx context.Context, account *Account) error {
  92. atomic.AddInt32(&r.updateCalled, 1)
  93. if r.updateErr != nil {
  94. return r.updateErr
  95. }
  96. r.account = account
  97. return nil
  98. }
  99. // openAIOAuthServiceStub implements OpenAIOAuthService methods for testing
  100. type openAIOAuthServiceStub struct {
  101. tokenInfo *OpenAITokenInfo
  102. refreshErr error
  103. refreshCalled int32
  104. }
  105. func (s *openAIOAuthServiceStub) RefreshAccountToken(ctx context.Context, account *Account) (*OpenAITokenInfo, error) {
  106. atomic.AddInt32(&s.refreshCalled, 1)
  107. if s.refreshErr != nil {
  108. return nil, s.refreshErr
  109. }
  110. return s.tokenInfo, nil
  111. }
  112. func (s *openAIOAuthServiceStub) BuildAccountCredentials(info *OpenAITokenInfo) map[string]any {
  113. now := time.Now()
  114. return map[string]any{
  115. "access_token": info.AccessToken,
  116. "refresh_token": info.RefreshToken,
  117. "expires_at": now.Add(time.Duration(info.ExpiresIn) * time.Second).Format(time.RFC3339),
  118. }
  119. }
  120. func TestOpenAITokenProvider_CacheHit(t *testing.T) {
  121. cache := newOpenAITokenCacheStub()
  122. account := &Account{
  123. ID: 100,
  124. Platform: PlatformOpenAI,
  125. Type: AccountTypeOAuth,
  126. Credentials: map[string]any{
  127. "access_token": "db-token",
  128. },
  129. }
  130. cacheKey := OpenAITokenCacheKey(account)
  131. cache.tokens[cacheKey] = "cached-token"
  132. provider := NewOpenAITokenProvider(nil, cache, nil)
  133. token, err := provider.GetAccessToken(context.Background(), account)
  134. require.NoError(t, err)
  135. require.Equal(t, "cached-token", token)
  136. require.Equal(t, int32(1), atomic.LoadInt32(&cache.getCalled))
  137. require.Equal(t, int32(0), atomic.LoadInt32(&cache.setCalled))
  138. }
  139. func TestOpenAITokenProvider_CacheMiss_FromCredentials(t *testing.T) {
  140. cache := newOpenAITokenCacheStub()
  141. // Token expires in far future, no refresh needed
  142. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  143. account := &Account{
  144. ID: 101,
  145. Platform: PlatformOpenAI,
  146. Type: AccountTypeOAuth,
  147. Credentials: map[string]any{
  148. "access_token": "credential-token",
  149. "expires_at": expiresAt,
  150. },
  151. }
  152. provider := NewOpenAITokenProvider(nil, cache, nil)
  153. token, err := provider.GetAccessToken(context.Background(), account)
  154. require.NoError(t, err)
  155. require.Equal(t, "credential-token", token)
  156. // Should have stored in cache
  157. cacheKey := OpenAITokenCacheKey(account)
  158. require.Equal(t, "credential-token", cache.tokens[cacheKey])
  159. }
  160. func TestOpenAITokenProvider_TokenRefresh(t *testing.T) {
  161. cache := newOpenAITokenCacheStub()
  162. accountRepo := &openAIAccountRepoStub{}
  163. oauthService := &openAIOAuthServiceStub{
  164. tokenInfo: &OpenAITokenInfo{
  165. AccessToken: "refreshed-token",
  166. RefreshToken: "new-refresh-token",
  167. ExpiresIn: 3600,
  168. },
  169. }
  170. // Token expires soon (within refresh skew)
  171. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  172. account := &Account{
  173. ID: 102,
  174. Platform: PlatformOpenAI,
  175. Type: AccountTypeOAuth,
  176. Credentials: map[string]any{
  177. "access_token": "old-token",
  178. "refresh_token": "old-refresh-token",
  179. "expires_at": expiresAt,
  180. },
  181. }
  182. accountRepo.account = account
  183. // We need to directly test with the stub - create a custom provider
  184. customProvider := &testOpenAITokenProvider{
  185. accountRepo: accountRepo,
  186. tokenCache: cache,
  187. oauthService: oauthService,
  188. }
  189. token, err := customProvider.GetAccessToken(context.Background(), account)
  190. require.NoError(t, err)
  191. require.Equal(t, "refreshed-token", token)
  192. require.Equal(t, int32(1), atomic.LoadInt32(&oauthService.refreshCalled))
  193. }
  194. // testOpenAITokenProvider is a test version that uses the stub OAuth service
  195. type testOpenAITokenProvider struct {
  196. accountRepo *openAIAccountRepoStub
  197. tokenCache *openAITokenCacheStub
  198. oauthService *openAIOAuthServiceStub
  199. }
  200. func (p *testOpenAITokenProvider) GetAccessToken(ctx context.Context, account *Account) (string, error) {
  201. if account == nil {
  202. return "", errors.New("account is nil")
  203. }
  204. if account.Platform != PlatformOpenAI || account.Type != AccountTypeOAuth {
  205. return "", errors.New("not an openai oauth account")
  206. }
  207. cacheKey := OpenAITokenCacheKey(account)
  208. // 1. Check cache
  209. if p.tokenCache != nil {
  210. if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && token != "" {
  211. return token, nil
  212. }
  213. }
  214. // 2. Check if refresh needed
  215. expiresAt := account.GetCredentialAsTime("expires_at")
  216. needsRefresh := expiresAt == nil || time.Until(*expiresAt) <= openAITokenRefreshSkew
  217. refreshFailed := false
  218. if needsRefresh && p.tokenCache != nil {
  219. locked, err := p.tokenCache.AcquireRefreshLock(ctx, cacheKey, 30*time.Second)
  220. if err == nil && locked {
  221. defer func() { _ = p.tokenCache.ReleaseRefreshLock(ctx, cacheKey) }()
  222. // Check cache again after acquiring lock
  223. if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && token != "" {
  224. return token, nil
  225. }
  226. // Get fresh account from DB
  227. fresh, err := p.accountRepo.GetByID(ctx, account.ID)
  228. if err == nil && fresh != nil {
  229. account = fresh
  230. }
  231. expiresAt = account.GetCredentialAsTime("expires_at")
  232. if expiresAt == nil || time.Until(*expiresAt) <= openAITokenRefreshSkew {
  233. if p.oauthService == nil {
  234. refreshFailed = true // 无法刷新,标记失败
  235. } else {
  236. tokenInfo, err := p.oauthService.RefreshAccountToken(ctx, account)
  237. if err != nil {
  238. refreshFailed = true // 刷新失败,标记以使用短 TTL
  239. } else {
  240. newCredentials := p.oauthService.BuildAccountCredentials(tokenInfo)
  241. for k, v := range account.Credentials {
  242. if _, exists := newCredentials[k]; !exists {
  243. newCredentials[k] = v
  244. }
  245. }
  246. account.Credentials = newCredentials
  247. _ = p.accountRepo.Update(ctx, account)
  248. expiresAt = account.GetCredentialAsTime("expires_at")
  249. }
  250. }
  251. }
  252. } else if p.tokenCache.simulateLockRace {
  253. // Wait and retry cache
  254. time.Sleep(10 * time.Millisecond) // Short wait for test
  255. if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && token != "" {
  256. return token, nil
  257. }
  258. }
  259. }
  260. accessToken := account.GetOpenAIAccessToken()
  261. if accessToken == "" {
  262. return "", errors.New("access_token not found in credentials")
  263. }
  264. // 3. Store in cache
  265. if p.tokenCache != nil {
  266. ttl := 30 * time.Minute
  267. if refreshFailed {
  268. ttl = time.Minute // 刷新失败时使用短 TTL
  269. } else if expiresAt != nil {
  270. until := time.Until(*expiresAt)
  271. if until > openAITokenCacheSkew {
  272. ttl = until - openAITokenCacheSkew
  273. } else if until > 0 {
  274. ttl = until
  275. } else {
  276. ttl = time.Minute
  277. }
  278. }
  279. _ = p.tokenCache.SetAccessToken(ctx, cacheKey, accessToken, ttl)
  280. }
  281. return accessToken, nil
  282. }
  283. func TestOpenAITokenProvider_LockRaceCondition(t *testing.T) {
  284. cache := newOpenAITokenCacheStub()
  285. cache.simulateLockRace = true
  286. accountRepo := &openAIAccountRepoStub{}
  287. // Token expires soon
  288. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  289. account := &Account{
  290. ID: 103,
  291. Platform: PlatformOpenAI,
  292. Type: AccountTypeOAuth,
  293. Credentials: map[string]any{
  294. "access_token": "race-token",
  295. "expires_at": expiresAt,
  296. },
  297. }
  298. accountRepo.account = account
  299. // Simulate another worker already refreshed and cached
  300. cacheKey := OpenAITokenCacheKey(account)
  301. go func() {
  302. time.Sleep(5 * time.Millisecond)
  303. cache.mu.Lock()
  304. cache.tokens[cacheKey] = "winner-token"
  305. cache.mu.Unlock()
  306. }()
  307. provider := &testOpenAITokenProvider{
  308. accountRepo: accountRepo,
  309. tokenCache: cache,
  310. }
  311. token, err := provider.GetAccessToken(context.Background(), account)
  312. require.NoError(t, err)
  313. // Should get the token set by the "winner" or the original
  314. require.NotEmpty(t, token)
  315. }
  316. func TestOpenAITokenProvider_NilAccount(t *testing.T) {
  317. provider := NewOpenAITokenProvider(nil, nil, nil)
  318. token, err := provider.GetAccessToken(context.Background(), nil)
  319. require.Error(t, err)
  320. require.Contains(t, err.Error(), "account is nil")
  321. require.Empty(t, token)
  322. }
  323. func TestOpenAITokenProvider_WrongPlatform(t *testing.T) {
  324. provider := NewOpenAITokenProvider(nil, nil, nil)
  325. account := &Account{
  326. ID: 104,
  327. Platform: PlatformGemini,
  328. Type: AccountTypeOAuth,
  329. }
  330. token, err := provider.GetAccessToken(context.Background(), account)
  331. require.Error(t, err)
  332. require.Contains(t, err.Error(), "not an openai oauth account")
  333. require.Empty(t, token)
  334. }
  335. func TestOpenAITokenProvider_WrongAccountType(t *testing.T) {
  336. provider := NewOpenAITokenProvider(nil, nil, nil)
  337. account := &Account{
  338. ID: 105,
  339. Platform: PlatformOpenAI,
  340. Type: AccountTypeAPIKey,
  341. }
  342. token, err := provider.GetAccessToken(context.Background(), account)
  343. require.Error(t, err)
  344. require.Contains(t, err.Error(), "not an openai oauth account")
  345. require.Empty(t, token)
  346. }
  347. func TestOpenAITokenProvider_NilCache(t *testing.T) {
  348. // Token doesn't need refresh
  349. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  350. account := &Account{
  351. ID: 106,
  352. Platform: PlatformOpenAI,
  353. Type: AccountTypeOAuth,
  354. Credentials: map[string]any{
  355. "access_token": "nocache-token",
  356. "expires_at": expiresAt,
  357. },
  358. }
  359. provider := NewOpenAITokenProvider(nil, nil, nil)
  360. token, err := provider.GetAccessToken(context.Background(), account)
  361. require.NoError(t, err)
  362. require.Equal(t, "nocache-token", token)
  363. }
  364. func TestOpenAITokenProvider_CacheGetError(t *testing.T) {
  365. cache := newOpenAITokenCacheStub()
  366. cache.getErr = errors.New("redis connection failed")
  367. // Token doesn't need refresh
  368. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  369. account := &Account{
  370. ID: 107,
  371. Platform: PlatformOpenAI,
  372. Type: AccountTypeOAuth,
  373. Credentials: map[string]any{
  374. "access_token": "fallback-token",
  375. "refresh_token": "refresh-token",
  376. "expires_at": expiresAt,
  377. },
  378. }
  379. provider := NewOpenAITokenProvider(nil, cache, nil)
  380. // Should gracefully degrade and return from credentials
  381. token, err := provider.GetAccessToken(context.Background(), account)
  382. require.NoError(t, err)
  383. require.Equal(t, "fallback-token", token)
  384. }
  385. func TestOpenAITokenProvider_CacheSetError(t *testing.T) {
  386. cache := newOpenAITokenCacheStub()
  387. cache.setErr = errors.New("redis write failed")
  388. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  389. account := &Account{
  390. ID: 108,
  391. Platform: PlatformOpenAI,
  392. Type: AccountTypeOAuth,
  393. Credentials: map[string]any{
  394. "access_token": "still-works-token",
  395. "expires_at": expiresAt,
  396. },
  397. }
  398. provider := NewOpenAITokenProvider(nil, cache, nil)
  399. // Should still work even if cache set fails
  400. token, err := provider.GetAccessToken(context.Background(), account)
  401. require.NoError(t, err)
  402. require.Equal(t, "still-works-token", token)
  403. }
  404. func TestOpenAITokenProvider_MissingAccessToken(t *testing.T) {
  405. cache := newOpenAITokenCacheStub()
  406. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  407. account := &Account{
  408. ID: 109,
  409. Platform: PlatformOpenAI,
  410. Type: AccountTypeOAuth,
  411. Credentials: map[string]any{
  412. "expires_at": expiresAt,
  413. // missing access_token
  414. },
  415. }
  416. provider := NewOpenAITokenProvider(nil, cache, nil)
  417. token, err := provider.GetAccessToken(context.Background(), account)
  418. require.Error(t, err)
  419. require.Contains(t, err.Error(), "access_token not found")
  420. require.Empty(t, token)
  421. }
  422. func TestOpenAITokenProvider_RefreshError(t *testing.T) {
  423. cache := newOpenAITokenCacheStub()
  424. accountRepo := &openAIAccountRepoStub{}
  425. oauthService := &openAIOAuthServiceStub{
  426. refreshErr: errors.New("oauth refresh failed"),
  427. }
  428. // Token expires soon
  429. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  430. account := &Account{
  431. ID: 110,
  432. Platform: PlatformOpenAI,
  433. Type: AccountTypeOAuth,
  434. Credentials: map[string]any{
  435. "access_token": "old-token",
  436. "refresh_token": "old-refresh-token",
  437. "expires_at": expiresAt,
  438. },
  439. }
  440. accountRepo.account = account
  441. provider := &testOpenAITokenProvider{
  442. accountRepo: accountRepo,
  443. tokenCache: cache,
  444. oauthService: oauthService,
  445. }
  446. // Now with fallback behavior, should return existing token even if refresh fails
  447. token, err := provider.GetAccessToken(context.Background(), account)
  448. require.NoError(t, err)
  449. require.Equal(t, "old-token", token) // Fallback to existing token
  450. }
  451. func TestOpenAITokenProvider_OAuthServiceNotConfigured(t *testing.T) {
  452. cache := newOpenAITokenCacheStub()
  453. accountRepo := &openAIAccountRepoStub{}
  454. // Token expires soon
  455. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  456. account := &Account{
  457. ID: 111,
  458. Platform: PlatformOpenAI,
  459. Type: AccountTypeOAuth,
  460. Credentials: map[string]any{
  461. "access_token": "old-token",
  462. "expires_at": expiresAt,
  463. },
  464. }
  465. accountRepo.account = account
  466. provider := &testOpenAITokenProvider{
  467. accountRepo: accountRepo,
  468. tokenCache: cache,
  469. oauthService: nil, // not configured
  470. }
  471. // Now with fallback behavior, should return existing token even if oauth service not configured
  472. token, err := provider.GetAccessToken(context.Background(), account)
  473. require.NoError(t, err)
  474. require.Equal(t, "old-token", token) // Fallback to existing token
  475. }
  476. func TestOpenAITokenProvider_TTLCalculation(t *testing.T) {
  477. tests := []struct {
  478. name string
  479. expiresIn time.Duration
  480. }{
  481. {
  482. name: "far_future_expiry",
  483. expiresIn: 1 * time.Hour,
  484. },
  485. {
  486. name: "medium_expiry",
  487. expiresIn: 10 * time.Minute,
  488. },
  489. {
  490. name: "near_expiry",
  491. expiresIn: 6 * time.Minute,
  492. },
  493. }
  494. for _, tt := range tests {
  495. t.Run(tt.name, func(t *testing.T) {
  496. cache := newOpenAITokenCacheStub()
  497. expiresAt := time.Now().Add(tt.expiresIn).Format(time.RFC3339)
  498. account := &Account{
  499. ID: 200,
  500. Platform: PlatformOpenAI,
  501. Type: AccountTypeOAuth,
  502. Credentials: map[string]any{
  503. "access_token": "test-token",
  504. "expires_at": expiresAt,
  505. },
  506. }
  507. provider := NewOpenAITokenProvider(nil, cache, nil)
  508. _, err := provider.GetAccessToken(context.Background(), account)
  509. require.NoError(t, err)
  510. // Verify token was cached
  511. cacheKey := OpenAITokenCacheKey(account)
  512. require.Equal(t, "test-token", cache.tokens[cacheKey])
  513. })
  514. }
  515. }
  516. func TestOpenAITokenProvider_DoubleCheckAfterLock(t *testing.T) {
  517. cache := newOpenAITokenCacheStub()
  518. accountRepo := &openAIAccountRepoStub{}
  519. oauthService := &openAIOAuthServiceStub{
  520. tokenInfo: &OpenAITokenInfo{
  521. AccessToken: "refreshed-token",
  522. RefreshToken: "new-refresh",
  523. ExpiresIn: 3600,
  524. },
  525. }
  526. // Token expires soon
  527. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  528. account := &Account{
  529. ID: 112,
  530. Platform: PlatformOpenAI,
  531. Type: AccountTypeOAuth,
  532. Credentials: map[string]any{
  533. "access_token": "old-token",
  534. "expires_at": expiresAt,
  535. },
  536. }
  537. accountRepo.account = account
  538. cacheKey := OpenAITokenCacheKey(account)
  539. // Simulate: first GetAccessToken returns empty, but after lock acquired, cache has token
  540. originalGet := int32(0)
  541. cache.tokens[cacheKey] = "" // Empty initially
  542. provider := &testOpenAITokenProvider{
  543. accountRepo: accountRepo,
  544. tokenCache: cache,
  545. oauthService: oauthService,
  546. }
  547. // In a goroutine, set the cached token after a small delay (simulating race)
  548. go func() {
  549. time.Sleep(5 * time.Millisecond)
  550. cache.mu.Lock()
  551. cache.tokens[cacheKey] = "cached-by-other"
  552. cache.mu.Unlock()
  553. }()
  554. token, err := provider.GetAccessToken(context.Background(), account)
  555. require.NoError(t, err)
  556. // Should get either the refreshed token or the cached one
  557. require.NotEmpty(t, token)
  558. _ = originalGet // Suppress unused warning
  559. }
  560. // Tests for real provider - to increase coverage
  561. func TestOpenAITokenProvider_Real_LockFailedWait(t *testing.T) {
  562. cache := newOpenAITokenCacheStub()
  563. cache.lockAcquired = false // Lock acquisition fails
  564. // Token expires soon (within refresh skew) to trigger lock attempt
  565. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  566. account := &Account{
  567. ID: 200,
  568. Platform: PlatformOpenAI,
  569. Type: AccountTypeOAuth,
  570. Credentials: map[string]any{
  571. "access_token": "fallback-token",
  572. "refresh_token": "refresh-token",
  573. "expires_at": expiresAt,
  574. },
  575. }
  576. // Set token in cache after lock wait period (simulate other worker refreshing)
  577. cacheKey := OpenAITokenCacheKey(account)
  578. go func() {
  579. time.Sleep(100 * time.Millisecond)
  580. cache.mu.Lock()
  581. cache.tokens[cacheKey] = "refreshed-by-other"
  582. cache.mu.Unlock()
  583. }()
  584. provider := NewOpenAITokenProvider(nil, cache, nil)
  585. token, err := provider.GetAccessToken(context.Background(), account)
  586. require.NoError(t, err)
  587. // Should get either the fallback token or the refreshed one
  588. require.NotEmpty(t, token)
  589. }
  590. func TestOpenAITokenProvider_Real_CacheHitAfterWait(t *testing.T) {
  591. cache := newOpenAITokenCacheStub()
  592. cache.lockAcquired = false // Lock acquisition fails
  593. // Token expires soon
  594. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  595. account := &Account{
  596. ID: 201,
  597. Platform: PlatformOpenAI,
  598. Type: AccountTypeOAuth,
  599. Credentials: map[string]any{
  600. "access_token": "original-token",
  601. "expires_at": expiresAt,
  602. },
  603. }
  604. cacheKey := OpenAITokenCacheKey(account)
  605. // Set token in cache immediately after wait starts
  606. go func() {
  607. time.Sleep(50 * time.Millisecond)
  608. cache.mu.Lock()
  609. cache.tokens[cacheKey] = "winner-token"
  610. cache.mu.Unlock()
  611. }()
  612. provider := NewOpenAITokenProvider(nil, cache, nil)
  613. token, err := provider.GetAccessToken(context.Background(), account)
  614. require.NoError(t, err)
  615. require.NotEmpty(t, token)
  616. }
  617. func TestOpenAITokenProvider_Real_ExpiredWithoutRefreshToken(t *testing.T) {
  618. cache := newOpenAITokenCacheStub()
  619. cache.lockAcquired = false // Prevent entering refresh logic
  620. // Token with nil expires_at (no expiry set) - should use credentials
  621. account := &Account{
  622. ID: 202,
  623. Platform: PlatformOpenAI,
  624. Type: AccountTypeOAuth,
  625. Credentials: map[string]any{
  626. "access_token": "no-expiry-token",
  627. },
  628. }
  629. provider := NewOpenAITokenProvider(nil, cache, nil)
  630. token, err := provider.GetAccessToken(context.Background(), account)
  631. // Without OAuth service, refresh will fail but token should be returned from credentials
  632. require.NoError(t, err)
  633. require.Equal(t, "no-expiry-token", token)
  634. }
  635. func TestOpenAITokenProvider_Real_WhitespaceToken(t *testing.T) {
  636. cache := newOpenAITokenCacheStub()
  637. cacheKey := "openai:account:203"
  638. cache.tokens[cacheKey] = " " // Whitespace only - should be treated as empty
  639. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  640. account := &Account{
  641. ID: 203,
  642. Platform: PlatformOpenAI,
  643. Type: AccountTypeOAuth,
  644. Credentials: map[string]any{
  645. "access_token": "real-token",
  646. "expires_at": expiresAt,
  647. },
  648. }
  649. provider := NewOpenAITokenProvider(nil, cache, nil)
  650. token, err := provider.GetAccessToken(context.Background(), account)
  651. require.NoError(t, err)
  652. require.Equal(t, "real-token", token) // Should fall back to credentials
  653. }
  654. func TestOpenAITokenProvider_Real_LockError(t *testing.T) {
  655. cache := newOpenAITokenCacheStub()
  656. cache.lockErr = errors.New("redis lock failed")
  657. // Token expires soon (within refresh skew)
  658. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  659. account := &Account{
  660. ID: 204,
  661. Platform: PlatformOpenAI,
  662. Type: AccountTypeOAuth,
  663. Credentials: map[string]any{
  664. "access_token": "fallback-on-lock-error",
  665. "expires_at": expiresAt,
  666. },
  667. }
  668. provider := NewOpenAITokenProvider(nil, cache, nil)
  669. token, err := provider.GetAccessToken(context.Background(), account)
  670. require.NoError(t, err)
  671. require.Equal(t, "fallback-on-lock-error", token)
  672. }
  673. func TestOpenAITokenProvider_Real_WhitespaceCredentialToken(t *testing.T) {
  674. cache := newOpenAITokenCacheStub()
  675. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  676. account := &Account{
  677. ID: 205,
  678. Platform: PlatformOpenAI,
  679. Type: AccountTypeOAuth,
  680. Credentials: map[string]any{
  681. "access_token": " ", // Whitespace only
  682. "expires_at": expiresAt,
  683. },
  684. }
  685. provider := NewOpenAITokenProvider(nil, cache, nil)
  686. token, err := provider.GetAccessToken(context.Background(), account)
  687. require.Error(t, err)
  688. require.Contains(t, err.Error(), "access_token not found")
  689. require.Empty(t, token)
  690. }
  691. func TestOpenAITokenProvider_Real_NilCredentials(t *testing.T) {
  692. cache := newOpenAITokenCacheStub()
  693. expiresAt := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
  694. account := &Account{
  695. ID: 206,
  696. Platform: PlatformOpenAI,
  697. Type: AccountTypeOAuth,
  698. Credentials: map[string]any{
  699. "expires_at": expiresAt,
  700. // No access_token
  701. },
  702. }
  703. provider := NewOpenAITokenProvider(nil, cache, nil)
  704. token, err := provider.GetAccessToken(context.Background(), account)
  705. require.Error(t, err)
  706. require.Contains(t, err.Error(), "access_token not found")
  707. require.Empty(t, token)
  708. }
  709. func TestOpenAITokenProvider_Real_LockRace_PollingHitsCache(t *testing.T) {
  710. cache := newOpenAITokenCacheStub()
  711. cache.lockAcquired = false // 模拟锁被其他 worker 持有
  712. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  713. account := &Account{
  714. ID: 207,
  715. Platform: PlatformOpenAI,
  716. Type: AccountTypeOAuth,
  717. Credentials: map[string]any{
  718. "access_token": "fallback-token",
  719. "refresh_token": "refresh-token",
  720. "expires_at": expiresAt,
  721. },
  722. }
  723. cacheKey := OpenAITokenCacheKey(account)
  724. go func() {
  725. time.Sleep(5 * time.Millisecond)
  726. cache.mu.Lock()
  727. cache.tokens[cacheKey] = "winner-token"
  728. cache.mu.Unlock()
  729. }()
  730. provider := NewOpenAITokenProvider(nil, cache, nil)
  731. token, err := provider.GetAccessToken(context.Background(), account)
  732. require.NoError(t, err)
  733. require.Equal(t, "winner-token", token)
  734. }
  735. func TestOpenAITokenProvider_Real_LockRace_ContextCanceled(t *testing.T) {
  736. cache := newOpenAITokenCacheStub()
  737. cache.lockAcquired = false // 模拟锁被其他 worker 持有
  738. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  739. account := &Account{
  740. ID: 208,
  741. Platform: PlatformOpenAI,
  742. Type: AccountTypeOAuth,
  743. Credentials: map[string]any{
  744. "access_token": "fallback-token",
  745. "refresh_token": "refresh-token",
  746. "expires_at": expiresAt,
  747. },
  748. }
  749. ctx, cancel := context.WithCancel(context.Background())
  750. cancel()
  751. provider := NewOpenAITokenProvider(nil, cache, nil)
  752. start := time.Now()
  753. token, err := provider.GetAccessToken(ctx, account)
  754. require.Error(t, err)
  755. require.ErrorIs(t, err, context.Canceled)
  756. require.Empty(t, token)
  757. require.Less(t, time.Since(start), 50*time.Millisecond)
  758. }
  759. func TestOpenAITokenProvider_RuntimeMetrics_LockWaitHitAndSnapshot(t *testing.T) {
  760. cache := newOpenAITokenCacheStub()
  761. cache.lockAcquired = false
  762. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  763. account := &Account{
  764. ID: 209,
  765. Platform: PlatformOpenAI,
  766. Type: AccountTypeOAuth,
  767. Credentials: map[string]any{
  768. "access_token": "fallback-token",
  769. "refresh_token": "refresh-token",
  770. "expires_at": expiresAt,
  771. },
  772. }
  773. cacheKey := OpenAITokenCacheKey(account)
  774. go func() {
  775. time.Sleep(10 * time.Millisecond)
  776. cache.mu.Lock()
  777. cache.tokens[cacheKey] = "winner-token"
  778. cache.mu.Unlock()
  779. }()
  780. provider := NewOpenAITokenProvider(nil, cache, nil)
  781. token, err := provider.GetAccessToken(context.Background(), account)
  782. require.NoError(t, err)
  783. require.Equal(t, "winner-token", token)
  784. metrics := provider.SnapshotRuntimeMetrics()
  785. require.GreaterOrEqual(t, metrics.RefreshRequests, int64(1))
  786. require.GreaterOrEqual(t, metrics.LockContention, int64(1))
  787. require.GreaterOrEqual(t, metrics.LockWaitSamples, int64(1))
  788. require.GreaterOrEqual(t, metrics.LockWaitHit, int64(1))
  789. require.GreaterOrEqual(t, metrics.LockWaitTotalMs, int64(0))
  790. require.GreaterOrEqual(t, metrics.LastObservedUnixMs, int64(1))
  791. }
  792. func TestOpenAITokenProvider_RuntimeMetrics_LockAcquireFailure(t *testing.T) {
  793. cache := newOpenAITokenCacheStub()
  794. cache.lockErr = errors.New("redis lock error")
  795. expiresAt := time.Now().Add(1 * time.Minute).Format(time.RFC3339)
  796. account := &Account{
  797. ID: 210,
  798. Platform: PlatformOpenAI,
  799. Type: AccountTypeOAuth,
  800. Credentials: map[string]any{
  801. "access_token": "fallback-token",
  802. "refresh_token": "refresh-token",
  803. "expires_at": expiresAt,
  804. },
  805. }
  806. provider := NewOpenAITokenProvider(nil, cache, nil)
  807. _, err := provider.GetAccessToken(context.Background(), account)
  808. require.NoError(t, err)
  809. metrics := provider.SnapshotRuntimeMetrics()
  810. require.GreaterOrEqual(t, metrics.LockAcquireFailure, int64(1))
  811. require.GreaterOrEqual(t, metrics.RefreshRequests, int64(1))
  812. }
  813. func TestOpenAITokenProvider_NoRefreshTokenExpired_DisablesAccount(t *testing.T) {
  814. cache := newOpenAITokenCacheStub()
  815. repo := &rateLimitAccountRepoStub{}
  816. expiresAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339)
  817. account := &Account{
  818. ID: 2881,
  819. Platform: PlatformOpenAI,
  820. Type: AccountTypeOAuth,
  821. Credentials: map[string]any{
  822. "access_token": "expired-access-token",
  823. "expires_at": expiresAt,
  824. },
  825. }
  826. cacheKey := OpenAITokenCacheKey(account)
  827. cache.tokens[cacheKey] = "stale-cached-token"
  828. // Force the provider past the cache hit branch.
  829. cache.getErr = errors.New("simulated cache miss")
  830. provider := NewOpenAITokenProvider(repo, cache, nil)
  831. blocker := &runtimeBlockRecorder{}
  832. provider.SetAccountRuntimeBlocker(blocker)
  833. token, err := provider.GetAccessToken(context.Background(), account)
  834. require.Error(t, err)
  835. require.Empty(t, token)
  836. require.Contains(t, err.Error(), "refresh_token is missing")
  837. require.Equal(t, 1, repo.setErrorCalls, "account should be disabled via SetError exactly once")
  838. require.Contains(t, repo.lastErrorMsg, "refresh_token is missing")
  839. require.Len(t, blocker.accounts, 1)
  840. require.Equal(t, account.ID, blocker.accounts[0].ID)
  841. require.Equal(t, "missing_refresh_token", blocker.reasons[0])
  842. }