package controller import ( "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/stretchr/testify/assert" "github.com/stretchr/testify/require" ) 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 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 TestDoubaoAssetProxyBlocksAssetGroupActionsCaseInsensitively(t *testing.T) { router := setupDoubaoAssetProxyRouter(t) for _, action := range []string{"CreateAssetGroup", "getassetgroup", "LISTASSETGROUPS", "UpdateAssetGroup", "deleteassetgroup"} { t.Run(action, func(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, nil) w := httptest.NewRecorder() router.ServeHTTP(w, req) assert.Equal(t, http.StatusBadRequest, w.Code) assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "Asset group API ("+action+") is not supported") }) } } func TestBuildDoubaoAssetURLDefaultBaseAndEscapedQuery(t *testing.T) { got, err := buildDoubaoAssetURL(&model.Channel{}, "ApplyUploadInner&Space", "") require.NoError(t, err) assert.Equal(t, defaultDoubaoAssetBaseURL+"/api/v1/multimodal/sd/assets?Action=ApplyUploadInner%26Space&Version=2024-01-01", got) } func TestBuildDoubaoAssetURLExplicitBaseURLOverridesDefault(t *testing.T) { baseURL := "https://example.com/custom/" got, err := buildDoubaoAssetURL(&model.Channel{BaseURL: &baseURL}, "CommitUploadInner", "2025-02-03") require.NoError(t, err) assert.Equal(t, "https://example.com/custom/api/v1/multimodal/sd/assets?Action=CommitUploadInner&Version=2025-02-03", got) } 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) }