account_long_context_billing_test.go 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "net/http"
  6. "testing"
  7. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestAccountIsOpenAILongContextBillingEnabled(t *testing.T) {
  11. tests := []struct {
  12. name string
  13. account *Account
  14. want bool
  15. }{
  16. {name: "nil account is disabled", account: nil, want: false},
  17. {name: "non OpenAI account is disabled", account: &Account{Platform: PlatformGrok}, want: false},
  18. {name: "missing extra defaults disabled", account: &Account{Platform: PlatformOpenAI}, want: false},
  19. {name: "missing key defaults disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{}}, want: false},
  20. {name: "explicit true is enabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": true}}, want: true},
  21. {name: "explicit false is disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": false}}, want: false},
  22. {name: "malformed value is disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": "false"}}, want: false},
  23. }
  24. for _, tt := range tests {
  25. t.Run(tt.name, func(t *testing.T) {
  26. require.Equal(t, tt.want, tt.account.IsOpenAILongContextBillingEnabled())
  27. })
  28. }
  29. }
  30. func TestNormalizeOpenAILongContextBillingExtra(t *testing.T) {
  31. t.Run("OpenAI missing key persists disabled default", func(t *testing.T) {
  32. extra, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, nil)
  33. require.NoError(t, err)
  34. require.Equal(t, false, extra["openai_long_context_billing_enabled"])
  35. })
  36. t.Run("OpenAI explicit false is preserved", func(t *testing.T) {
  37. extra, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, map[string]any{"openai_long_context_billing_enabled": false})
  38. require.NoError(t, err)
  39. require.Equal(t, false, extra["openai_long_context_billing_enabled"])
  40. })
  41. t.Run("OpenAI malformed value is rejected", func(t *testing.T) {
  42. _, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, map[string]any{"openai_long_context_billing_enabled": "false"})
  43. require.Error(t, err)
  44. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  45. })
  46. t.Run("non OpenAI extra is unchanged", func(t *testing.T) {
  47. extra, err := normalizeOpenAILongContextBillingExtra(PlatformGrok, nil)
  48. require.NoError(t, err)
  49. require.Nil(t, extra)
  50. })
  51. t.Run("non OpenAI malformed value is ignored", func(t *testing.T) {
  52. extra := map[string]any{openAILongContextBillingEnabledKey: "provider-owned"}
  53. normalized, err := normalizeOpenAILongContextBillingExtra(PlatformAnthropic, extra)
  54. require.NoError(t, err)
  55. require.Equal(t, extra, normalized)
  56. })
  57. }
  58. type longContextBillingRepoStub struct {
  59. accountRepoStub
  60. account *Account
  61. accounts []*Account
  62. createdAccount *Account
  63. updateExtraCalls int
  64. bulkUpdateCalls int
  65. }
  66. func (r *longContextBillingRepoStub) Create(_ context.Context, account *Account) error {
  67. account.ID = 1
  68. r.account = account
  69. r.createdAccount = account
  70. return nil
  71. }
  72. func (r *longContextBillingRepoStub) GetByID(_ context.Context, _ int64) (*Account, error) {
  73. return r.account, nil
  74. }
  75. func (r *longContextBillingRepoStub) GetByIDs(_ context.Context, _ []int64) ([]*Account, error) {
  76. if r.accounts != nil {
  77. return r.accounts, nil
  78. }
  79. if r.account == nil {
  80. return nil, nil
  81. }
  82. return []*Account{r.account}, nil
  83. }
  84. func (r *longContextBillingRepoStub) Update(_ context.Context, account *Account) error {
  85. r.account = account
  86. return nil
  87. }
  88. func (r *longContextBillingRepoStub) UpdateExtra(_ context.Context, _ int64, _ map[string]any) error {
  89. r.updateExtraCalls++
  90. return nil
  91. }
  92. func (r *longContextBillingRepoStub) BulkUpdate(_ context.Context, _ []int64, _ AccountBulkUpdate) (int64, error) {
  93. r.bulkUpdateCalls++
  94. return 1, nil
  95. }
  96. func TestAdminServiceCreateAccountDefaultsOpenAILongContextBillingDisabled(t *testing.T) {
  97. repo := &longContextBillingRepoStub{}
  98. svc := &adminServiceImpl{accountRepo: repo}
  99. account, err := svc.CreateAccount(context.Background(), &CreateAccountInput{
  100. Name: "openai-account",
  101. Platform: PlatformOpenAI,
  102. Type: AccountTypeAPIKey,
  103. Credentials: map[string]any{"api_key": "test"},
  104. SkipDefaultGroupBind: true,
  105. })
  106. require.NoError(t, err)
  107. require.Same(t, account, repo.createdAccount)
  108. require.Equal(t, false, account.Extra[openAILongContextBillingEnabledKey])
  109. }
  110. func TestAdminServiceCreateAccountRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
  111. repo := &longContextBillingRepoStub{}
  112. svc := &adminServiceImpl{accountRepo: repo}
  113. account, err := svc.CreateAccount(context.Background(), &CreateAccountInput{
  114. Platform: PlatformOpenAI,
  115. Extra: map[string]any{openAILongContextBillingEnabledKey: "false"},
  116. })
  117. require.Nil(t, account)
  118. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  119. require.Nil(t, repo.createdAccount)
  120. }
  121. func TestAdminServiceUpdateAccountPreservesOpenAILongContextBillingOptOutWhenOmitted(t *testing.T) {
  122. repo := &longContextBillingRepoStub{account: &Account{
  123. ID: 1,
  124. Platform: PlatformOpenAI,
  125. Type: AccountTypeOAuth,
  126. Extra: map[string]any{openAILongContextBillingEnabledKey: false},
  127. }}
  128. svc := &adminServiceImpl{accountRepo: repo}
  129. account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{}})
  130. require.NoError(t, err)
  131. require.Equal(t, false, account.Extra[openAILongContextBillingEnabledKey])
  132. }
  133. func TestAdminServiceUpdateAccountAllowsExplicitCodexImportOptIn(t *testing.T) {
  134. repo := &longContextBillingRepoStub{account: &Account{
  135. ID: 1,
  136. Platform: PlatformOpenAI,
  137. Type: AccountTypeOAuth,
  138. Credentials: map[string]any{"access_token": "old-token"},
  139. Extra: map[string]any{
  140. openAILongContextBillingEnabledKey: false,
  141. "import_source": "codex_session",
  142. },
  143. }}
  144. svc := &adminServiceImpl{accountRepo: repo}
  145. account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{
  146. Credentials: map[string]any{"access_token": "new-token"},
  147. Extra: map[string]any{
  148. openAILongContextBillingEnabledKey: true,
  149. "import_source": "codex_session",
  150. },
  151. })
  152. require.NoError(t, err)
  153. require.Equal(t, true, account.Extra[openAILongContextBillingEnabledKey])
  154. }
  155. func TestAdminServiceUpdateAccountAllowsExplicitOptInOutsideCodexImport(t *testing.T) {
  156. repo := &longContextBillingRepoStub{account: &Account{
  157. ID: 1,
  158. Platform: PlatformOpenAI,
  159. Type: AccountTypeOAuth,
  160. Extra: map[string]any{
  161. openAILongContextBillingEnabledKey: false,
  162. "import_source": "codex_session",
  163. },
  164. }}
  165. svc := &adminServiceImpl{accountRepo: repo}
  166. account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{
  167. openAILongContextBillingEnabledKey: true,
  168. "import_source": "codex_session",
  169. }})
  170. require.NoError(t, err)
  171. require.Equal(t, true, account.Extra[openAILongContextBillingEnabledKey])
  172. }
  173. func TestAdminServiceUpdateAccountRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
  174. repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}}
  175. svc := &adminServiceImpl{accountRepo: repo}
  176. account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{
  177. openAILongContextBillingEnabledKey: 1,
  178. }})
  179. require.Nil(t, account)
  180. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  181. }
  182. func TestAdminServiceUpdateAccountExtraRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
  183. repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}}
  184. svc := &adminServiceImpl{accountRepo: repo}
  185. err := svc.UpdateAccountExtra(context.Background(), 1, map[string]any{
  186. openAILongContextBillingEnabledKey: "true",
  187. })
  188. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  189. require.Zero(t, repo.updateExtraCalls)
  190. }
  191. func TestAdminServiceUpdateAccountExtraAllowsProviderOwnedValueForNonOpenAIAccount(t *testing.T) {
  192. repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformAnthropic}}
  193. svc := &adminServiceImpl{accountRepo: repo}
  194. err := svc.UpdateAccountExtra(context.Background(), 1, map[string]any{
  195. openAILongContextBillingEnabledKey: "provider-owned",
  196. })
  197. require.NoError(t, err)
  198. require.Equal(t, 1, repo.updateExtraCalls)
  199. }
  200. func TestAdminServiceBulkUpdateAccountsRejectsMalformedOpenAILongContextBillingValue(t *testing.T) {
  201. repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}}
  202. svc := &adminServiceImpl{accountRepo: repo}
  203. result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
  204. AccountIDs: []int64{1},
  205. Extra: map[string]any{openAILongContextBillingEnabledKey: []bool{true}},
  206. })
  207. require.Nil(t, result)
  208. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  209. require.Zero(t, repo.bulkUpdateCalls)
  210. }
  211. func TestAdminServiceBulkUpdateAccountsRejectsOpenAILongContextKeyForNonOpenAIAccounts(t *testing.T) {
  212. repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformGrok}}
  213. svc := &adminServiceImpl{accountRepo: repo}
  214. result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
  215. AccountIDs: []int64{1},
  216. Extra: map[string]any{openAILongContextBillingEnabledKey: true},
  217. })
  218. require.Nil(t, result)
  219. var appErr *infraerrors.ApplicationError
  220. require.ErrorAs(t, err, &appErr)
  221. require.Equal(t, "OPENAI_BULK_TARGET_INVALID", appErr.Reason)
  222. require.Zero(t, repo.bulkUpdateCalls)
  223. }
  224. func TestAdminServiceBulkUpdateAccountsRejectsMalformedValueForMixedTargetsIncludingOpenAI(t *testing.T) {
  225. repo := &longContextBillingRepoStub{accounts: []*Account{
  226. {ID: 1, Platform: PlatformGrok},
  227. {ID: 2, Platform: PlatformOpenAI},
  228. }}
  229. svc := &adminServiceImpl{accountRepo: repo}
  230. result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
  231. AccountIDs: []int64{1, 2},
  232. Extra: map[string]any{openAILongContextBillingEnabledKey: "malformed"},
  233. })
  234. require.Nil(t, result)
  235. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  236. require.Zero(t, repo.bulkUpdateCalls)
  237. }