openai_privacy_retry_test.go 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "github.com/Wei-Shaw/sub2api/internal/config"
  8. "github.com/imroc/req/v3"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestAdminService_EnsureOpenAIPrivacy_RetriesNonSuccessModes(t *testing.T) {
  12. t.Parallel()
  13. for _, mode := range []string{PrivacyModeFailed, PrivacyModeCFBlocked} {
  14. t.Run(mode, func(t *testing.T) {
  15. t.Parallel()
  16. privacyCalls := 0
  17. svc := &adminServiceImpl{
  18. accountRepo: &mockAccountRepoForGemini{},
  19. privacyClientFactory: func(proxyURL string) (*req.Client, error) {
  20. privacyCalls++
  21. return nil, errors.New("factory failed")
  22. },
  23. }
  24. account := &Account{
  25. ID: 101,
  26. Platform: PlatformOpenAI,
  27. Type: AccountTypeOAuth,
  28. Credentials: map[string]any{
  29. "access_token": "token-1",
  30. },
  31. Extra: map[string]any{
  32. "privacy_mode": mode,
  33. },
  34. }
  35. got := svc.EnsureOpenAIPrivacy(context.Background(), account)
  36. require.Equal(t, PrivacyModeFailed, got)
  37. require.Equal(t, 1, privacyCalls)
  38. })
  39. }
  40. }
  41. func TestTokenRefreshService_ensureOpenAIPrivacy_RetriesNonSuccessModes(t *testing.T) {
  42. t.Parallel()
  43. cfg := &config.Config{
  44. TokenRefresh: config.TokenRefreshConfig{
  45. MaxRetries: 1,
  46. RetryBackoffSeconds: 0,
  47. },
  48. }
  49. for _, mode := range []string{PrivacyModeFailed, PrivacyModeCFBlocked} {
  50. t.Run(mode, func(t *testing.T) {
  51. t.Parallel()
  52. service := NewTokenRefreshService(&tokenRefreshAccountRepo{}, nil, nil, nil, nil, nil, nil, cfg, nil)
  53. privacyCalls := 0
  54. service.SetPrivacyDeps(func(proxyURL string) (*req.Client, error) {
  55. privacyCalls++
  56. return nil, errors.New("factory failed")
  57. }, nil)
  58. account := &Account{
  59. ID: 202,
  60. Platform: PlatformOpenAI,
  61. Type: AccountTypeOAuth,
  62. Credentials: map[string]any{
  63. "access_token": "token-2",
  64. },
  65. Extra: map[string]any{
  66. "privacy_mode": mode,
  67. },
  68. }
  69. service.ensureOpenAIPrivacy(context.Background(), account)
  70. require.Equal(t, 1, privacyCalls)
  71. })
  72. }
  73. }