package controller import ( "bytes" "net/http" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func setupCodexChannelDB(t *testing.T, key string) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) sqlDB, _ := db.DB() sqlDB.SetMaxOpenConns(1) origDB := model.DB origLogDB := model.LOG_DB model.DB = db model.LOG_DB = db common.UsingSQLite = true common.RedisEnabled = false require.NoError(t, db.AutoMigrate(&model.Channel{})) require.NoError(t, db.Create(&model.Channel{ Id: 1, Name: "codex-channel", PublicName: "codex-channel", Type: constant.ChannelTypeCodex, Key: key, Status: common.ChannelStatusEnabled, }).Error) t.Cleanup(func() { model.DB = origDB model.LOG_DB = origLogDB _ = sqlDB.Close() }) return db } func setupCodexChannelRouter(t *testing.T, key string) *gin.Engine { t.Helper() setupCodexChannelDB(t, key) gin.SetMode(gin.TestMode) r := gin.New() g := r.Group("/api/channel") g.POST("/:id/codex/refresh", RefreshCodexChannelCredential) g.GET("/:id/codex/usage", GetCodexChannelUsage) return r } func TestValidateChannelAcceptsCodexAPIKey(t *testing.T) { channel := &model.Channel{ Name: "codex-api-key", PublicName: "codex-api-key", Type: constant.ChannelTypeCodex, Key: "sk-codex-api-key", } require.NoError(t, validateChannel(channel, true)) } func TestValidateChannelRejectsCodexOAuthWithoutAccountID(t *testing.T) { channel := &model.Channel{ Name: "codex-oauth", PublicName: "codex-oauth", Type: constant.ChannelTypeCodex, Key: `{"access_token":"token-only"}`, } err := validateChannel(channel, true) require.Error(t, err) assert.Contains(t, err.Error(), "account_id") } func TestRefreshCodexChannelCredentialRejectsAPIKeyMode(t *testing.T) { router := setupCodexChannelRouter(t, "sk-codex-api-key") req := httptest.NewRequest(http.MethodPost, "/api/channel/1/codex/refresh", bytes.NewReader([]byte(`{}`))) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() router.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) assert.Contains(t, w.Body.String(), "当前凭证方式不支持刷新凭证") } func TestGetCodexChannelUsageRejectsAPIKeyMode(t *testing.T) { router := setupCodexChannelRouter(t, "sk-codex-api-key") req := httptest.NewRequest(http.MethodGet, "/api/channel/1/codex/usage", nil) w := httptest.NewRecorder() router.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) assert.Contains(t, w.Body.String(), "当前凭证方式不支持查看用量") }