new-api/controller/doubao_asset_real_e2e_test.go

package controller

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"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/require"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

const chinaMobileAssetE2EImageURL = "https://bkimg.cdn.bcebos.com/pic/caef76094b36acaf2edd2133a78e9a1001e9380136fe?x-bce-process=image/format,f_auto/watermark,image_d2F0ZXIvYmFpa2UyNzI,g_7,xp_5,yp_5,P_20/resize,m_lfit,limit_1,h_1080"

func TestChinaMobileAssetControllerWithRealUpstream(t *testing.T) {
	ak := strings.TrimSpace(os.Getenv("CHINAMOBILE_ASSET_AK"))
	sk := strings.TrimSpace(os.Getenv("CHINAMOBILE_ASSET_SK"))
	poolID := strings.TrimSpace(os.Getenv("CHINAMOBILE_ASSET_POOL_ID"))
	if ak == "" || sk == "" {
		t.Skip("CHINAMOBILE_ASSET_AK and CHINAMOBILE_ASSET_SK are required")
	}
	if poolID == "" {
		poolID = "CIDC-CORE-00"
	}
	db := setupRealAssetE2EDB(t)
	createDoubaoAssetProxyChannel(t, db, 7101, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
	router := setupRealAssetE2ERouter()
	suffix := time.Now().Format("20060102150405")
	groupName := "new-api-e2e-" + suffix
	assetName := "new-api-img-" + suffix

	var groupId string
	var assetId string
	defer func() {
		if assetId != "" {
			code, _ := performRealAssetE2EAction(t, router, "DeleteAsset", map[string]any{"Id": assetId})
			require.Equal(t, http.StatusOK, code, "cleanup DeleteAsset failed")
		}
		if groupId != "" {
			code, _ := performRealAssetE2EAction(t, router, "DeleteAssetGroup", map[string]any{"Id": groupId})
			require.Equal(t, http.StatusOK, code, "cleanup DeleteAssetGroup failed")
		}
	}()

	code, body := performRealAssetE2EAction(t, router, "CreateAssetGroup", map[string]any{
		"Name":        groupName,
		"GroupType":   "AIGC",
		"Description": "new-api real e2e test",
	})
	require.Equal(t, http.StatusOK, code)
	groupId = realAssetE2EString(t, body, "Result.GroupId")
	require.NotEmpty(t, groupId)

	code, body = performRealAssetE2EAction(t, router, "CreateAsset", map[string]any{
		"GroupId":   groupId,
		"Name":      assetName,
		"URL":       chinaMobileAssetE2EImageURL,
		"AssetType": "Image",
	})
	require.Equal(t, http.StatusOK, code)
	assetId = realAssetE2EString(t, body, "Result")
	require.NotEmpty(t, assetId)

	var status string
	for i := 0; i < 6; i++ {
		code, body = performRealAssetE2EAction(t, router, "GetAsset", map[string]any{"Id": assetId})
		require.Equal(t, http.StatusOK, code)
		require.Equal(t, assetId, realAssetE2EString(t, body, "Result.Id"))
		status = realAssetE2EString(t, body, "Result.Status")
		if status != "Processing" {
			break
		}
		time.Sleep(3 * time.Second)
	}
	require.Equal(t, "Active", status)

	code, body = performRealAssetE2EAction(t, router, "ListAssets", map[string]any{
		"PageNumber": 1,
		"PageSize":   10,
		"Filter": map[string]any{
			"GroupType": "AIGC",
			"GroupIds":  []string{groupId},
		},
	})
	require.Equal(t, http.StatusOK, code)
	require.Equal(t, float64(1), realAssetE2EValue(t, body, "Result.TotalCount"))
}

func setupRealAssetE2EDB(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{}))

	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 setupRealAssetE2ERouter() *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 performRealAssetE2EAction(t *testing.T, router *gin.Engine, action string, payload map[string]any) (int, string) {
	t.Helper()
	data, err := common.Marshal(payload)
	require.NoError(t, err)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, bytes.NewReader(data))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	return w.Code, w.Body.String()
}

func realAssetE2EString(t *testing.T, data string, path string) string {
	t.Helper()
	value := realAssetE2EValue(t, data, path)
	s, ok := value.(string)
	require.Truef(t, ok, "%s is not a string: %#v", path, value)
	return s
}

func realAssetE2EValue(t *testing.T, data string, path string) any {
	t.Helper()
	var payload map[string]any
	require.NoError(t, common.Unmarshal([]byte(data), &payload))
	current := any(payload)
	for _, part := range strings.Split(path, ".") {
		m, ok := current.(map[string]any)
		require.Truef(t, ok, "%s is not an object at %s", path, part)
		current = m[part]
	}
	return current
}