Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.
 
 
 

112 righe
2.8 KiB

  1. package controller
  2. import (
  3. "bytes"
  4. "net/http"
  5. "net/http/httptest"
  6. "testing"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/constant"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/gin-gonic/gin"
  11. "github.com/glebarez/sqlite"
  12. "github.com/stretchr/testify/assert"
  13. "github.com/stretchr/testify/require"
  14. "gorm.io/gorm"
  15. )
  16. func setupCodexChannelDB(t *testing.T, key string) *gorm.DB {
  17. t.Helper()
  18. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  19. require.NoError(t, err)
  20. sqlDB, _ := db.DB()
  21. sqlDB.SetMaxOpenConns(1)
  22. origDB := model.DB
  23. origLogDB := model.LOG_DB
  24. model.DB = db
  25. model.LOG_DB = db
  26. common.UsingSQLite = true
  27. common.RedisEnabled = false
  28. require.NoError(t, db.AutoMigrate(&model.Channel{}))
  29. require.NoError(t, db.Create(&model.Channel{
  30. Id: 1,
  31. Name: "codex-channel",
  32. PublicName: "codex-channel",
  33. Type: constant.ChannelTypeCodex,
  34. Key: key,
  35. Status: common.ChannelStatusEnabled,
  36. }).Error)
  37. t.Cleanup(func() {
  38. model.DB = origDB
  39. model.LOG_DB = origLogDB
  40. _ = sqlDB.Close()
  41. })
  42. return db
  43. }
  44. func setupCodexChannelRouter(t *testing.T, key string) *gin.Engine {
  45. t.Helper()
  46. setupCodexChannelDB(t, key)
  47. gin.SetMode(gin.TestMode)
  48. r := gin.New()
  49. g := r.Group("/api/channel")
  50. g.POST("/:id/codex/refresh", RefreshCodexChannelCredential)
  51. g.GET("/:id/codex/usage", GetCodexChannelUsage)
  52. return r
  53. }
  54. func TestValidateChannelAcceptsCodexAPIKey(t *testing.T) {
  55. channel := &model.Channel{
  56. Name: "codex-api-key",
  57. PublicName: "codex-api-key",
  58. Type: constant.ChannelTypeCodex,
  59. Key: "sk-codex-api-key",
  60. }
  61. require.NoError(t, validateChannel(channel, true))
  62. }
  63. func TestValidateChannelRejectsCodexOAuthWithoutAccountID(t *testing.T) {
  64. channel := &model.Channel{
  65. Name: "codex-oauth",
  66. PublicName: "codex-oauth",
  67. Type: constant.ChannelTypeCodex,
  68. Key: `{"access_token":"token-only"}`,
  69. }
  70. err := validateChannel(channel, true)
  71. require.Error(t, err)
  72. assert.Contains(t, err.Error(), "account_id")
  73. }
  74. func TestRefreshCodexChannelCredentialRejectsAPIKeyMode(t *testing.T) {
  75. router := setupCodexChannelRouter(t, "sk-codex-api-key")
  76. req := httptest.NewRequest(http.MethodPost, "/api/channel/1/codex/refresh", bytes.NewReader([]byte(`{}`)))
  77. req.Header.Set("Content-Type", "application/json")
  78. w := httptest.NewRecorder()
  79. router.ServeHTTP(w, req)
  80. assert.Equal(t, http.StatusOK, w.Code)
  81. assert.Contains(t, w.Body.String(), "当前凭证方式不支持刷新凭证")
  82. }
  83. func TestGetCodexChannelUsageRejectsAPIKeyMode(t *testing.T) {
  84. router := setupCodexChannelRouter(t, "sk-codex-api-key")
  85. req := httptest.NewRequest(http.MethodGet, "/api/channel/1/codex/usage", nil)
  86. w := httptest.NewRecorder()
  87. router.ServeHTTP(w, req)
  88. assert.Equal(t, http.StatusOK, w.Code)
  89. assert.Contains(t, w.Body.String(), "当前凭证方式不支持查看用量")
  90. }