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 }