|
- package controller
-
- import (
- "context"
- "fmt"
- "net/http"
- "net/http/httptest"
- "os"
- "strings"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/constant"
- "github.com/QuantumNous/new-api/model"
- "github.com/QuantumNous/new-api/service"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- "gorm.io/gorm/logger"
- )
-
- func setupDoubaoAssetProxyRouter(t *testing.T) *gin.Engine {
- t.Helper()
-
- oldMode := gin.Mode()
- gin.SetMode(gin.TestMode)
- t.Cleanup(func() {
- gin.SetMode(oldMode)
- })
-
- r := gin.New()
- r.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
- return r
- }
-
- func decodeDoubaoAssetErrorMessage(t *testing.T, body string) string {
- t.Helper()
-
- var payload struct {
- Error struct {
- Message string `json:"message"`
- Type string `json:"type"`
- } `json:"error"`
- }
- require.NoError(t, common.Unmarshal([]byte(body), &payload))
- return payload.Error.Message
- }
-
- func setupDoubaoAssetProxyDB(t *testing.T) *gorm.DB {
- t.Helper()
-
- oldDB := model.DB
- oldSQLitePath := common.SQLitePath
- oldMemoryCacheEnabled := common.MemoryCacheEnabled
- oldIsMasterNode := common.IsMasterNode
- oldUsingSQLite := common.UsingSQLite
- oldUsingMySQL := common.UsingMySQL
- oldUsingPostgreSQL := common.UsingPostgreSQL
- oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")
-
- common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared"
- common.MemoryCacheEnabled = false
- common.IsMasterNode = false
- common.UsingSQLite = false
- common.UsingMySQL = false
- common.UsingPostgreSQL = false
- require.NoError(t, os.Setenv("SQL_DSN", "local"))
- require.NoError(t, model.InitDB())
-
- model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
- db := model.DB
- sqlDB, err := db.DB()
- require.NoError(t, err)
- sqlDB.SetMaxOpenConns(1)
- require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}, &model.UserAssetGroup{}))
-
- t.Cleanup(func() {
- _ = sqlDB.Close()
- model.DB = oldDB
- common.SQLitePath = oldSQLitePath
- common.MemoryCacheEnabled = oldMemoryCacheEnabled
- common.IsMasterNode = oldIsMasterNode
- common.UsingSQLite = oldUsingSQLite
- common.UsingMySQL = oldUsingMySQL
- common.UsingPostgreSQL = oldUsingPostgreSQL
- if hadSQLDSN {
- _ = os.Setenv("SQL_DSN", oldSQLDSN)
- } else {
- _ = os.Unsetenv("SQL_DSN")
- }
- })
-
- return db
- }
-
- func createDoubaoAssetProxyChannel(t *testing.T, db *gorm.DB, id int, channelType int, group string, key string, status int) {
- t.Helper()
-
- priority := int64(id)
- weight := uint(10)
- autoBan := 1
- require.NoError(t, db.Create(&model.Channel{
- Id: id,
- Type: channelType,
- Key: key,
- Status: status,
- Name: fmt.Sprintf("channel-%d", id),
- Group: group,
- Models: "seedance-2",
- Priority: &priority,
- Weight: &weight,
- AutoBan: &autoBan,
- CreatedTime: int64(id),
- }).Error)
- }
-
- func TestDoubaoAssetProxyMissingActionReturns400(t *testing.T) {
- router := setupDoubaoAssetProxyRouter(t)
-
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Version=2024-01-01", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusBadRequest, w.Code)
- assert.Equal(t, "Action query parameter is required", decodeDoubaoAssetErrorMessage(t, w.Body.String()))
- }
-
- func TestDoubaoAssetProxyUnknownActionReturns400(t *testing.T) {
- router := setupDoubaoAssetProxyRouter(t)
-
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=CreateRealPersonAuthSession", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusBadRequest, w.Code)
- assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "unsupported asset Action")
- }
-
- func TestDoubaoAssetProxyInvalidJSONReturns400(t *testing.T) {
- router := setupDoubaoAssetProxyRouter(t)
-
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"PageNumber":`))
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusBadRequest, w.Code)
- assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "invalid JSON body")
- }
-
- func TestDoubaoAssetProxyRejectsChinaMobileAssetGroupActions(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- fake := &fakeDoubaoAssetAdapter{}
- t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
- createDoubaoAssetProxyChannel(t, db, 61, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
- router := gin.New()
- router.Use(func(c *gin.Context) {
- c.Set("id", 10)
- common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
- c.Next()
- })
- router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
-
- for _, action := range []string{"CreateAssetGroup", "ListAssetGroups", "GetAssetGroup", "UpdateAssetGroup", "DeleteAssetGroup"} {
- t.Run(action, func(t *testing.T) {
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, strings.NewReader(`{"Id":"group-1","GroupType":"AIGC","Name":"g"}`))
- w := httptest.NewRecorder()
- req.Header.Set("Content-Type", "application/json")
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusForbidden, w.Code)
- assert.Equal(t, 0, fake.calls)
- })
- }
- }
-
- func TestDoubaoAssetProxyDoesNotBlockNonChinaMobileAssetGroupAction(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- fake := &fakeDoubaoAssetAdapter{}
- t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeDoubaoVideoCompatibleAiping, fake))
- createDoubaoAssetProxyChannel(t, db, 63, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "video-generation-key", common.ChannelStatusEnabled)
- router := newDoubaoAssetProxyTestRouter()
-
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssetGroups", strings.NewReader(`{"PageNumber":1,"PageSize":10}`))
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- assert.Equal(t, service.AssetOperationAssetGroupList, fake.operation)
- }
-
- func TestDoubaoAssetProxyAutoGroupContinuesAfterMissingGroup(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- fake := &fakeDoubaoAssetAdapter{}
- t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
- createDoubaoAssetProxyChannel(t, db, 62, constant.ChannelTypeChinaMobileSeedance, "vip", "video-generation-key", common.ChannelStatusEnabled)
- router := gin.New()
- router.Use(func(c *gin.Context) {
- c.Set("id", 10)
- common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
- common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
- c.Next()
- })
- router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
- oldAutoGroups := doubaoAssetAutoGroups
- doubaoAssetAutoGroups = func(string) []string {
- return []string{"default", "vip"}
- }
- t.Cleanup(func() {
- doubaoAssetAutoGroups = oldAutoGroups
- })
-
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
- w := httptest.NewRecorder()
- req.Header.Set("Content-Type", "application/json")
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- assert.Equal(t, service.AssetOperationAssetList, fake.operation)
- }
-
- func TestDoubaoAssetProxyForwardsAdapterResponseHeaders(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- fake := &fakeDoubaoAssetAdapter{responseHeader: http.Header{"X-Upstream-Request-Id": []string{"upstream-1"}, "Content-Length": []string{"999"}}}
- t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
- createDoubaoAssetProxyChannel(t, db, 62, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
- router := gin.New()
- router.Use(func(c *gin.Context) {
- c.Set("id", 10)
- common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
- c.Next()
- })
- router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
-
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
- w := httptest.NewRecorder()
- req.Header.Set("Content-Type", "application/json")
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusOK, w.Code)
- assert.Equal(t, "upstream-1", w.Header().Get("X-Upstream-Request-Id"))
- assert.NotEqual(t, "999", w.Header().Get("Content-Length"))
- }
-
- func TestDoubaoAssetProxyScopesChinaMobileCreateAndListRequests(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- fake := &fakeDoubaoAssetAdapter{}
- t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
- createDoubaoAssetProxyChannel(t, db, 71, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
- require.NoError(t, model.CreateUserAssetGroup(10, 71, "owned"))
- router := newDoubaoAssetProxyTestRouter()
-
- for _, tc := range []struct {
- action string
- body string
- }{
- {"CreateAsset", `{"GroupId":"forged","Name":"n","URL":"https://example.com/a.png","AssetType":"Image"}`},
- {"ListAssets", `{"Filter":{"GroupType":"AIGC","GroupIds":["forged"]},"PageNumber":1,"PageSize":10}`},
- } {
- t.Run(tc.action, func(t *testing.T) {
- fake.requests = nil
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+tc.action, strings.NewReader(tc.body))
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
- require.Equal(t, http.StatusOK, w.Code)
- require.Len(t, fake.requests, 1)
- if tc.action == "CreateAsset" {
- assert.Equal(t, "owned", fake.requests[0].Body["GroupId"])
- } else {
- assert.Equal(t, []string{"owned"}, fake.requests[0].Body["Filter"].(map[string]any)["GroupIds"])
- }
- })
- }
- }
-
- func TestDoubaoAssetProxyRejectsChinaMobileListAssetsWithoutFilter(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- createDoubaoAssetProxyChannel(t, db, 73, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
- require.NoError(t, model.CreateUserAssetGroup(10, 73, "owned"))
- router := newDoubaoAssetProxyTestRouter()
-
- for _, body := range []string{`{"PageNumber":1,"PageSize":10}`, `{"Filter":null,"PageNumber":1,"PageSize":10}`} {
- w := httptest.NewRecorder()
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(body))
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusBadRequest, w.Code)
- assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "Filter.GroupType is required")
- }
- }
-
- func TestDoubaoAssetProxyHidesChinaMobileAssetsOwnedByAnotherUser(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- fake := &fakeDoubaoAssetAdapter{getAssetGroupID: "other"}
- t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
- createDoubaoAssetProxyChannel(t, db, 72, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
- require.NoError(t, model.CreateUserAssetGroup(10, 72, "owned"))
- router := newDoubaoAssetProxyTestRouter()
-
- for _, action := range []string{"GetAsset", "UpdateAsset", "DeleteAsset"} {
- t.Run(action, func(t *testing.T) {
- fake.requests = nil
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, strings.NewReader(`{"Id":"asset-other","Name":"n"}`))
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusNotFound, w.Code)
- require.Len(t, fake.requests, 1)
- assert.Equal(t, service.AssetOperationAssetGet, fake.requests[0].Action.Operation)
- })
- }
- }
-
- func newDoubaoAssetProxyTestRouter() *gin.Engine {
- router := gin.New()
- router.Use(func(c *gin.Context) {
- c.Set("id", 10)
- common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
- c.Next()
- })
- router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
- return router
- }
-
- func TestDoubaoAssetProxyAutoGroupDoesNotBypassInvalidBinding(t *testing.T) {
- db := setupDoubaoAssetProxyDB(t)
- createDoubaoAssetProxyChannel(t, db, 1, constant.ChannelTypeChinaMobileSeedance, "other", "bad-key", common.ChannelStatusEnabled)
- createDoubaoAssetProxyChannel(t, db, 2, constant.ChannelTypeChinaMobileSeedance, "vip", "video-generation-key", common.ChannelStatusEnabled)
- require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeChinaMobileSeedance, "default", 1))
- router := gin.New()
- router.Use(func(c *gin.Context) {
- c.Set("id", 10)
- common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
- common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
- c.Next()
- })
- router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
- oldAutoGroups := doubaoAssetAutoGroups
- doubaoAssetAutoGroups = func(string) []string {
- return []string{"default", "vip"}
- }
- t.Cleanup(func() {
- doubaoAssetAutoGroups = oldAutoGroups
- })
-
- req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
- w := httptest.NewRecorder()
- req.Header.Set("Content-Type", "application/json")
- router.ServeHTTP(w, req)
-
- assert.Equal(t, http.StatusBadGateway, w.Code)
- assert.Contains(t, w.Body.String(), service.AssetErrorBindingInvalid)
- }
-
- func TestEffectiveDoubaoAssetGroupUsesUsingGroupBeforeBlankTokenGroup(t *testing.T) {
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- common.SetContextKey(c, constant.ContextKeyTokenGroup, "")
- common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
-
- assert.Equal(t, "default", effectiveDoubaoAssetGroup(c))
- }
-
- func TestConcreteDoubaoAssetGroupsForAutoUsesUserAutoGroups(t *testing.T) {
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
- common.SetContextKey(c, constant.ContextKeyTokenGroup, "")
-
- groups := concreteDoubaoAssetGroupsForRequest(c, func(string) []string {
- return []string{"default", "vip"}
- })
-
- assert.Equal(t, []string{"default", "vip"}, groups)
- }
-
- type fakeDoubaoAssetAdapter struct {
- operation service.AssetOperation
- responseHeader http.Header
- requests []service.AssetRequest
- getAssetGroupID string
- calls int
- }
-
- func (a *fakeDoubaoAssetAdapter) Name() string {
- return "fake_asset"
- }
-
- func (a *fakeDoubaoAssetAdapter) Supports(operation service.AssetOperation) bool {
- return true
- }
-
- func (a *fakeDoubaoAssetAdapter) DoAssetRequest(_ context.Context, _ *model.Channel, req service.AssetRequest) (*service.AssetUpstreamResponse, *service.AssetError) {
- a.operation = req.Action.Operation
- a.requests = append(a.requests, req)
- a.calls++
- if req.Action.Operation == service.AssetOperationAssetGroupCreate {
- body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{"GroupId": "generated-group"})
- if err != nil {
- return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
- }
- return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
- }
- if req.Action.Operation == service.AssetOperationAssetGet {
- body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{"Id": "asset-1", "GroupId": a.getAssetGroupID})
- if err != nil {
- return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
- }
- return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
- }
- body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{
- "Items": []map[string]any{
- {"Id": "group-1", "Name": "g", "GroupType": "AIGC"},
- },
- "TotalCount": 1,
- })
- if err != nil {
- return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
- }
- return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
- }
|