openai_codex_pat_service_test.go 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. package service
  2. import (
  3. "context"
  4. "net/http"
  5. "net/http/httptest"
  6. "testing"
  7. "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestOpenAIOAuthService_ValidateCodexPersonalAccessToken(t *testing.T) {
  11. var gotAuthorization string
  12. var gotOriginator string
  13. var gotUserAgent string
  14. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  15. gotAuthorization = r.Header.Get("authorization")
  16. gotOriginator = r.Header.Get("originator")
  17. gotUserAgent = r.Header.Get("user-agent")
  18. w.Header().Set("content-type", "application/json")
  19. _, _ = w.Write([]byte(`{
  20. "email":"user@example.com",
  21. "chatgpt_user_id":"user-123",
  22. "chatgpt_account_id":"acct-123",
  23. "chatgpt_plan_type":"plus",
  24. "chatgpt_account_is_fedramp":true
  25. }`))
  26. }))
  27. defer server.Close()
  28. originalURL := openAICodexPATWhoamiURL
  29. openAICodexPATWhoamiURL = server.URL
  30. defer func() { openAICodexPATWhoamiURL = originalURL }()
  31. svc := NewOpenAIOAuthService(nil, nil)
  32. defer svc.Stop()
  33. info, err := svc.ValidateCodexPersonalAccessToken(context.Background(), " at-test-token ", "")
  34. require.NoError(t, err)
  35. require.Equal(t, "Bearer at-test-token", gotAuthorization)
  36. require.Equal(t, openai.CodexDefaultOriginator, gotOriginator)
  37. require.Equal(t, CodexCanonicalUserAgent(), gotUserAgent)
  38. require.Equal(t, OpenAIAuthModePersonalAccessToken, info.AuthMode)
  39. require.Equal(t, "user@example.com", info.Email)
  40. require.Equal(t, "user-123", info.ChatGPTUserID)
  41. require.Equal(t, "acct-123", info.ChatGPTAccountID)
  42. require.Equal(t, "plus", info.PlanType)
  43. require.True(t, info.ChatGPTAccountFedRAMP)
  44. require.Zero(t, info.ExpiresAt)
  45. require.Empty(t, info.RefreshToken)
  46. }
  47. func TestOpenAIOAuthService_ValidateCodexPersonalAccessTokenRequiresATPrefix(t *testing.T) {
  48. svc := NewOpenAIOAuthService(nil, nil)
  49. defer svc.Stop()
  50. _, err := svc.ValidateCodexPersonalAccessToken(context.Background(), "eyJ.jwt", "")
  51. require.Error(t, err)
  52. require.Contains(t, err.Error(), "at-")
  53. }
  54. func TestOpenAIOAuthService_BuildAccountCredentialsForPAT(t *testing.T) {
  55. svc := NewOpenAIOAuthService(nil, nil)
  56. defer svc.Stop()
  57. credentials := svc.BuildAccountCredentials(&OpenAITokenInfo{
  58. AccessToken: "at-test-token",
  59. AuthMode: OpenAIAuthModePersonalAccessToken,
  60. Email: "user@example.com",
  61. ChatGPTAccountID: "acct-123",
  62. ChatGPTUserID: "user-123",
  63. ChatGPTAccountFedRAMP: true,
  64. PlanType: "plus",
  65. })
  66. require.Equal(t, "at-test-token", credentials["access_token"])
  67. require.Equal(t, OpenAIAuthModePersonalAccessToken, credentials["auth_mode"])
  68. require.Equal(t, "personal_access_token", credentials["openai_auth_mode"])
  69. require.Equal(t, "Bearer", credentials["token_type"])
  70. require.Equal(t, true, credentials["chatgpt_account_is_fedramp"])
  71. require.NotContains(t, credentials, "expires_at")
  72. require.NotContains(t, credentials, "refresh_token")
  73. require.NotContains(t, credentials, "id_token")
  74. }
  75. func TestNormalizeOpenAIPersonalAccessTokenCredentialsRemovesOAuthFields(t *testing.T) {
  76. account := &Account{
  77. Platform: PlatformOpenAI,
  78. Type: AccountTypeOAuth,
  79. Credentials: map[string]any{
  80. "auth_mode": "personal_access_token",
  81. },
  82. }
  83. credentials := map[string]any{
  84. "access_token": "at-test-token",
  85. "refresh_token": "stale-refresh-token",
  86. "id_token": "stale-id-token",
  87. "expires_at": "2026-01-01T00:00:00Z",
  88. "expires_in": 3600,
  89. "client_id": "stale-client",
  90. "model_mapping": map[string]any{"gpt-5": "gpt-5-codex"},
  91. "chatgpt_account_is_fedramp": true,
  92. "subscription_expires_at": "2026-12-31T00:00:00Z",
  93. "openai_usage_channel_fields": []any{"custom"},
  94. }
  95. got := NormalizeOpenAIPersonalAccessTokenCredentials(account, nil, credentials)
  96. require.Equal(t, "at-test-token", got["access_token"])
  97. require.Equal(t, OpenAIAuthModePersonalAccessToken, got["auth_mode"])
  98. require.Equal(t, "personal_access_token", got["openai_auth_mode"])
  99. require.Equal(t, "Bearer", got["token_type"])
  100. require.NotContains(t, got, "refresh_token")
  101. require.NotContains(t, got, "id_token")
  102. require.NotContains(t, got, "expires_at")
  103. require.NotContains(t, got, "expires_in")
  104. require.NotContains(t, got, "client_id")
  105. require.Equal(t, map[string]any{"gpt-5": "gpt-5-codex"}, got["model_mapping"])
  106. require.Equal(t, true, got["chatgpt_account_is_fedramp"])
  107. require.Equal(t, "2026-12-31T00:00:00Z", got["subscription_expires_at"])
  108. require.Equal(t, []any{"custom"}, got["openai_usage_channel_fields"])
  109. }