new-api/controller/doubao_asset_test.go

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)
}