account_handler_long_context_billing_test.go 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165
  1. package admin
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  9. "github.com/Wei-Shaw/sub2api/internal/service"
  10. "github.com/gin-gonic/gin"
  11. "github.com/stretchr/testify/require"
  12. )
  13. func TestAccountAdminBoundariesRejectMalformedOpenAILongContextBillingValue(t *testing.T) {
  14. const malformedExtra = `"extra":{"openai_long_context_billing_enabled":"true"}`
  15. tests := []struct {
  16. name string
  17. method string
  18. path string
  19. body string
  20. mount func(*gin.Engine, *AccountHandler)
  21. setup func(*stubAdminService)
  22. }{
  23. {
  24. name: "create",
  25. method: http.MethodPost,
  26. path: "/accounts",
  27. body: `{"name":"account","platform":"openai","type":"apikey","credentials":{"api_key":"test"},` + malformedExtra + `}`,
  28. mount: func(router *gin.Engine, handler *AccountHandler) { router.POST("/accounts", handler.Create) },
  29. },
  30. {
  31. name: "update",
  32. method: http.MethodPut,
  33. path: "/accounts/1",
  34. body: `{` + malformedExtra + `}`,
  35. mount: func(router *gin.Engine, handler *AccountHandler) { router.PUT("/accounts/:id", handler.Update) },
  36. setup: func(stub *stubAdminService) {
  37. stub.updateAccountErr = infraerrors.BadRequest("OPENAI_LONG_CONTEXT_BILLING_INVALID", "invalid")
  38. },
  39. },
  40. {
  41. name: "bulk update",
  42. method: http.MethodPost,
  43. path: "/accounts/bulk-update",
  44. body: `{"account_ids":[1],` + malformedExtra + `}`,
  45. mount: func(router *gin.Engine, handler *AccountHandler) {
  46. router.POST("/accounts/bulk-update", handler.BulkUpdate)
  47. },
  48. setup: func(stub *stubAdminService) {
  49. stub.bulkUpdateAccountErr = infraerrors.BadRequest("OPENAI_LONG_CONTEXT_BILLING_INVALID", "invalid")
  50. },
  51. },
  52. {
  53. name: "batch create",
  54. method: http.MethodPost,
  55. path: "/accounts/batch",
  56. body: `{"accounts":[{"name":"account","platform":"openai","type":"apikey","credentials":{"api_key":"test"},` + malformedExtra + `}]}`,
  57. mount: func(router *gin.Engine, handler *AccountHandler) { router.POST("/accounts/batch", handler.BatchCreate) },
  58. },
  59. {
  60. name: "Codex session import",
  61. method: http.MethodPost,
  62. path: "/accounts/import-codex-session",
  63. body: `{"content":"token",` + malformedExtra + `}`,
  64. mount: func(router *gin.Engine, handler *AccountHandler) {
  65. router.POST("/accounts/import-codex-session", handler.ImportCodexSession)
  66. },
  67. },
  68. }
  69. for _, tt := range tests {
  70. t.Run(tt.name, func(t *testing.T) {
  71. gin.SetMode(gin.TestMode)
  72. stub := newStubAdminService()
  73. if tt.setup != nil {
  74. tt.setup(stub)
  75. }
  76. handler := NewAccountHandler(stub, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
  77. router := gin.New()
  78. tt.mount(router, handler)
  79. recorder := httptest.NewRecorder()
  80. request := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body))
  81. request.Header.Set("Content-Type", "application/json")
  82. router.ServeHTTP(recorder, request)
  83. require.Equal(t, http.StatusBadRequest, recorder.Code)
  84. var responseBody struct {
  85. Reason string `json:"reason"`
  86. }
  87. require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
  88. require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
  89. })
  90. }
  91. }
  92. func TestAccountCreateBoundaryDoesNotApplyOpenAIValidationToOtherPlatforms(t *testing.T) {
  93. gin.SetMode(gin.TestMode)
  94. handler := NewAccountHandler(newStubAdminService(), nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
  95. router := gin.New()
  96. router.POST("/accounts", handler.Create)
  97. recorder := httptest.NewRecorder()
  98. request := httptest.NewRequest(http.MethodPost, "/accounts", bytes.NewBufferString(
  99. `{"name":"account","platform":"anthropic","type":"apikey","credentials":{"api_key":"test"},"extra":{"openai_long_context_billing_enabled":"provider-owned"}}`,
  100. ))
  101. request.Header.Set("Content-Type", "application/json")
  102. router.ServeHTTP(recorder, request)
  103. require.Equal(t, http.StatusOK, recorder.Code)
  104. }
  105. func TestApplyOAuthCredentialsRejectsMalformedOpenAILongContextBillingBeforeMutation(t *testing.T) {
  106. gin.SetMode(gin.TestMode)
  107. stub := newStubAdminService()
  108. stub.getAccountResult = &service.Account{
  109. ID: 1,
  110. Platform: service.PlatformOpenAI,
  111. Type: service.AccountTypeOAuth,
  112. }
  113. handler := NewAccountHandler(stub, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
  114. router := gin.New()
  115. router.POST("/accounts/:id/apply-oauth-credentials", handler.ApplyOAuthCredentials)
  116. recorder := httptest.NewRecorder()
  117. request := httptest.NewRequest(http.MethodPost, "/accounts/1/apply-oauth-credentials", bytes.NewBufferString(
  118. `{"type":"oauth","credentials":{"access_token":"new-token"},"extra":{"openai_long_context_billing_enabled":"true"}}`,
  119. ))
  120. request.Header.Set("Content-Type", "application/json")
  121. router.ServeHTTP(recorder, request)
  122. require.Equal(t, http.StatusBadRequest, recorder.Code)
  123. var responseBody struct {
  124. Reason string `json:"reason"`
  125. }
  126. require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
  127. require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
  128. require.Zero(t, stub.updateAccountCalls)
  129. require.Zero(t, stub.updateAccountExtraCalls)
  130. }
  131. func TestOpenAIOAuthCodexPATBoundaryRejectsMalformedOpenAILongContextBillingValueBeforeTokenValidation(t *testing.T) {
  132. gin.SetMode(gin.TestMode)
  133. handler := NewOpenAIOAuthHandler(nil, newStubAdminService(), nil, nil)
  134. router := gin.New()
  135. router.Use(gin.Recovery())
  136. router.POST("/openai/create-from-codex-pat", handler.CreateAccountFromCodexPAT)
  137. recorder := httptest.NewRecorder()
  138. request := httptest.NewRequest(http.MethodPost, "/openai/create-from-codex-pat", bytes.NewBufferString(
  139. `{"access_token":"token","extra":{"openai_long_context_billing_enabled":1}}`,
  140. ))
  141. request.Header.Set("Content-Type", "application/json")
  142. router.ServeHTTP(recorder, request)
  143. require.Equal(t, http.StatusBadRequest, recorder.Code)
  144. var responseBody struct {
  145. Reason string `json:"reason"`
  146. }
  147. require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
  148. require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
  149. }