|
- 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(), "当前凭证方式不支持查看用量")
- }
|