| package admin |
|
|
| import ( |
| "context" |
| "encoding/json" |
| "io" |
| "net/http" |
| "net/http/httptest" |
| "strings" |
| "testing" |
|
|
| "github.com/Wei-Shaw/sub2api/internal/config" |
| "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" |
| "github.com/Wei-Shaw/sub2api/internal/service" |
| "github.com/gin-gonic/gin" |
| "github.com/stretchr/testify/require" |
| ) |
|
|
| type availableModelsAdminService struct { |
| *stubAdminService |
| account service.Account |
| } |
|
|
| func (s *availableModelsAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) { |
| if s.account.ID == id { |
| acc := s.account |
| return &acc, nil |
| } |
| return s.stubAdminService.GetAccount(context.Background(), id) |
| } |
|
|
| func setupAvailableModelsRouter(adminSvc service.AdminService) *gin.Engine { |
| gin.SetMode(gin.TestMode) |
| router := gin.New() |
| handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) |
| router.GET("/api/v1/admin/accounts/:id/models", handler.GetAvailableModels) |
| return router |
| } |
|
|
| type syncUpstreamHTTPUpstream struct { |
| resp *http.Response |
| err error |
| } |
|
|
| func (u *syncUpstreamHTTPUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { |
| if u.err != nil { |
| return nil, u.err |
| } |
| return u.resp, nil |
| } |
|
|
| func (u *syncUpstreamHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) { |
| return u.Do(req, proxyURL, accountID, accountConcurrency) |
| } |
|
|
| func setupSyncUpstreamModelsRouter(adminSvc service.AdminService, upstream service.HTTPUpstream) *gin.Engine { |
| gin.SetMode(gin.TestMode) |
| router := gin.New() |
| accountTestSvc := service.NewAccountTestService( |
| nil, |
| nil, |
| nil, |
| nil, |
| upstream, |
| &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}}, |
| nil, |
| ) |
| handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, accountTestSvc, nil, nil, nil, nil, nil) |
| router.POST("/api/v1/admin/accounts/:id/models/sync-upstream", handler.SyncUpstreamModels) |
| return router |
| } |
|
|
| func TestAccountHandlerGetAvailableModels_OpenAIOAuthUsesExplicitModelMapping(t *testing.T) { |
| svc := &availableModelsAdminService{ |
| stubAdminService: newStubAdminService(), |
| account: service.Account{ |
| ID: 42, |
| Name: "openai-oauth", |
| Platform: service.PlatformOpenAI, |
| Type: service.AccountTypeOAuth, |
| Status: service.StatusActive, |
| Credentials: map[string]any{ |
| "model_mapping": map[string]any{ |
| "gpt-5": "gpt-5.1", |
| }, |
| }, |
| }, |
| } |
| router := setupAvailableModelsRouter(svc) |
|
|
| rec := httptest.NewRecorder() |
| req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/42/models", nil) |
| router.ServeHTTP(rec, req) |
|
|
| require.Equal(t, http.StatusOK, rec.Code) |
|
|
| var resp struct { |
| Data []struct { |
| ID string `json:"id"` |
| } `json:"data"` |
| } |
| require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) |
| require.Len(t, resp.Data, 1) |
| require.Equal(t, "gpt-5", resp.Data[0].ID) |
| } |
|
|
| func TestAccountHandlerGetAvailableModels_OpenAIOAuthPassthroughFallsBackToDefaults(t *testing.T) { |
| svc := &availableModelsAdminService{ |
| stubAdminService: newStubAdminService(), |
| account: service.Account{ |
| ID: 43, |
| Name: "openai-oauth-passthrough", |
| Platform: service.PlatformOpenAI, |
| Type: service.AccountTypeOAuth, |
| Status: service.StatusActive, |
| Credentials: map[string]any{ |
| "model_mapping": map[string]any{ |
| "gpt-5": "gpt-5.1", |
| }, |
| }, |
| Extra: map[string]any{ |
| "openai_passthrough": true, |
| }, |
| }, |
| } |
| router := setupAvailableModelsRouter(svc) |
|
|
| rec := httptest.NewRecorder() |
| req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/43/models", nil) |
| router.ServeHTTP(rec, req) |
|
|
| require.Equal(t, http.StatusOK, rec.Code) |
|
|
| var resp struct { |
| Data []struct { |
| ID string `json:"id"` |
| } `json:"data"` |
| } |
| require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) |
| require.NotEmpty(t, resp.Data) |
| require.NotEqual(t, "gpt-5", resp.Data[0].ID) |
| } |
|
|
| func TestAccountHandlerSyncUpstreamModels_ConfigErrorReturnsBadRequest(t *testing.T) { |
| svc := &availableModelsAdminService{ |
| stubAdminService: newStubAdminService(), |
| account: service.Account{ |
| ID: 44, |
| Name: "openai-apikey-missing-key", |
| Platform: service.PlatformOpenAI, |
| Type: service.AccountTypeAPIKey, |
| Status: service.StatusActive, |
| Credentials: map[string]any{ |
| "base_url": "https://openai.example.com/v1", |
| }, |
| }, |
| } |
| router := setupSyncUpstreamModelsRouter(svc, &syncUpstreamHTTPUpstream{}) |
|
|
| rec := httptest.NewRecorder() |
| req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/44/models/sync-upstream", nil) |
| router.ServeHTTP(rec, req) |
|
|
| require.Equal(t, http.StatusBadRequest, rec.Code) |
| require.Contains(t, rec.Body.String(), "No OpenAI API key is available") |
| } |
|
|
| func TestAccountHandlerSyncUpstreamModels_UpstreamErrorDoesNotExposeBody(t *testing.T) { |
| svc := &availableModelsAdminService{ |
| stubAdminService: newStubAdminService(), |
| account: service.Account{ |
| ID: 45, |
| Name: "openai-apikey-upstream-error", |
| Platform: service.PlatformOpenAI, |
| Type: service.AccountTypeAPIKey, |
| Status: service.StatusActive, |
| Credentials: map[string]any{ |
| "api_key": "openai-key", |
| "base_url": "https://openai.example.com/v1", |
| }, |
| }, |
| } |
| upstream := &syncUpstreamHTTPUpstream{resp: &http.Response{ |
| StatusCode: http.StatusBadGateway, |
| Header: http.Header{"Content-Type": []string{"application/json"}}, |
| Body: io.NopCloser(strings.NewReader(`{"error":"SECRET_TOKEN should not be exposed"}`)), |
| }} |
| router := setupSyncUpstreamModelsRouter(svc, upstream) |
|
|
| rec := httptest.NewRecorder() |
| req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/45/models/sync-upstream", nil) |
| router.ServeHTTP(rec, req) |
|
|
| require.Equal(t, http.StatusBadGateway, rec.Code) |
| require.Contains(t, rec.Body.String(), "Upstream model list request failed with HTTP 502") |
| require.NotContains(t, rec.Body.String(), "SECRET_TOKEN") |
| } |
|
|