crs_sync_long_context_billing_test.go 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "encoding/json"
  6. "net/http"
  7. "net/http/httptest"
  8. "testing"
  9. "github.com/Wei-Shaw/sub2api/internal/config"
  10. "github.com/stretchr/testify/require"
  11. )
  12. type crsLongContextAccountRepo struct {
  13. AccountRepository
  14. accounts map[string]*Account
  15. nextID int64
  16. }
  17. type crsOpenAILongContextSource struct {
  18. collection string
  19. credentials map[string]any
  20. extra map[string]any
  21. }
  22. func newCRSLongContextAccountRepo(existing ...*Account) *crsLongContextAccountRepo {
  23. repo := &crsLongContextAccountRepo{accounts: make(map[string]*Account)}
  24. for _, account := range existing {
  25. if account == nil {
  26. continue
  27. }
  28. crsID, _ := account.Extra["crs_account_id"].(string)
  29. repo.accounts[crsID] = account
  30. if account.ID > repo.nextID {
  31. repo.nextID = account.ID
  32. }
  33. }
  34. return repo
  35. }
  36. func (r *crsLongContextAccountRepo) Create(_ context.Context, account *Account) error {
  37. r.nextID++
  38. account.ID = r.nextID
  39. crsID, _ := account.Extra["crs_account_id"].(string)
  40. r.accounts[crsID] = account
  41. return nil
  42. }
  43. func (r *crsLongContextAccountRepo) Update(_ context.Context, account *Account) error {
  44. crsID, _ := account.Extra["crs_account_id"].(string)
  45. r.accounts[crsID] = account
  46. return nil
  47. }
  48. func (r *crsLongContextAccountRepo) GetByCRSAccountID(_ context.Context, crsID string) (*Account, error) {
  49. return r.accounts[crsID], nil
  50. }
  51. func (r *crsLongContextAccountRepo) ListShadowsByParent(_ context.Context, _ int64) ([]*Account, error) {
  52. return nil, nil
  53. }
  54. func TestCRSSyncOpenAILongContextBilling(t *testing.T) {
  55. tests := []struct {
  56. name string
  57. collection string
  58. credentials map[string]any
  59. sourceExtra map[string]any
  60. existingExtra map[string]any
  61. wantAction string
  62. wantEnabled bool
  63. }{
  64. {name: "OAuth create defaults missing value disabled", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, wantAction: "created"},
  65. {name: "OAuth create preserves source true", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "created", wantEnabled: true},
  66. {name: "OAuth create preserves source false", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "created"},
  67. {name: "OAuth update defaults missing value disabled", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{"existing": true}, wantAction: "updated"},
  68. {name: "OAuth update preserves existing true when source omits value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated", wantEnabled: true},
  69. {name: "OAuth update preserves existing false when source omits value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated"},
  70. {name: "OAuth update preserves source true over existing false", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated", wantEnabled: true},
  71. {name: "OAuth update preserves source false over existing true", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated"},
  72. {name: "OAuth rejects malformed source value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"},
  73. {name: "OAuth rejects malformed existing value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"},
  74. {name: "OAuth update rejects malformed source value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "failed"},
  75. {name: "API key create defaults missing value disabled", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, wantAction: "created"},
  76. {name: "API key create preserves source true", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "created", wantEnabled: true},
  77. {name: "API key create preserves source false", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "created"},
  78. {name: "API key update defaults missing value disabled", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{"existing": true}, wantAction: "updated"},
  79. {name: "API key update preserves existing true when source omits value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated", wantEnabled: true},
  80. {name: "API key update preserves existing false when source omits value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated"},
  81. {name: "API key update preserves source true over existing false", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated", wantEnabled: true},
  82. {name: "API key update preserves source false over existing true", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated"},
  83. {name: "API key rejects malformed source value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"},
  84. {name: "API key rejects malformed existing value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"},
  85. {name: "API key update rejects malformed source value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "failed"},
  86. }
  87. for _, tt := range tests {
  88. t.Run(tt.name, func(t *testing.T) {
  89. const crsID = "crs-openai-1"
  90. var existing *Account
  91. if tt.existingExtra != nil {
  92. existingExtra := mergeMap(tt.existingExtra, map[string]any{"crs_account_id": crsID})
  93. accountType := AccountTypeOAuth
  94. if tt.collection == "openaiResponsesAccounts" {
  95. accountType = AccountTypeAPIKey
  96. }
  97. existing = &Account{ID: 41, Platform: PlatformOpenAI, Type: accountType, Extra: existingExtra}
  98. }
  99. repo := newCRSLongContextAccountRepo(existing)
  100. result := runCRSOpenAILongContextSync(t, repo, crsOpenAILongContextSource{
  101. collection: tt.collection,
  102. credentials: tt.credentials,
  103. extra: tt.sourceExtra,
  104. })
  105. require.Len(t, result.Items, 1)
  106. require.Equal(t, tt.wantAction, result.Items[0].Action)
  107. if tt.wantAction == "failed" {
  108. require.Contains(t, result.Items[0].Error, "openai_long_context_billing_enabled must be a boolean")
  109. return
  110. }
  111. stored, ok := repo.accounts[crsID].Extra[openAILongContextBillingEnabledKey]
  112. require.True(t, ok)
  113. require.Equal(t, tt.wantEnabled, stored)
  114. })
  115. }
  116. }
  117. func runCRSOpenAILongContextSync(t *testing.T, repo AccountRepository, source crsOpenAILongContextSource) *SyncFromCRSResult {
  118. t.Helper()
  119. account := map[string]any{
  120. "kind": "openai",
  121. "id": "crs-openai-1",
  122. "name": "OpenAI CRS",
  123. "isActive": true,
  124. "schedulable": true,
  125. "credentials": source.credentials,
  126. }
  127. if source.extra != nil {
  128. account["extra"] = source.extra
  129. }
  130. server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
  131. response.Header().Set("Content-Type", "application/json")
  132. if request.URL.Path == "/web/auth/login" {
  133. _, _ = response.Write([]byte(`{"success":true,"token":"admin-token"}`))
  134. return
  135. }
  136. require.Equal(t, "/admin/sync/export-accounts", request.URL.Path)
  137. require.NoError(t, json.NewEncoder(response).Encode(map[string]any{
  138. "success": true,
  139. "data": map[string]any{source.collection: []any{account}},
  140. }))
  141. }))
  142. t.Cleanup(server.Close)
  143. cfg := &config.Config{}
  144. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  145. service := NewCRSSyncService(repo, nil, nil, nil, nil, cfg)
  146. result, err := service.SyncFromCRS(context.Background(), SyncFromCRSInput{
  147. BaseURL: server.URL,
  148. Username: "admin",
  149. Password: "password",
  150. })
  151. require.NoError(t, err)
  152. return result
  153. }