Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

420 Zeilen
16 KiB

  1. package controller
  2. import (
  3. "context"
  4. "fmt"
  5. "net/http"
  6. "net/http/httptest"
  7. "os"
  8. "strings"
  9. "testing"
  10. "github.com/QuantumNous/new-api/common"
  11. "github.com/QuantumNous/new-api/constant"
  12. "github.com/QuantumNous/new-api/model"
  13. "github.com/QuantumNous/new-api/service"
  14. "github.com/gin-gonic/gin"
  15. "github.com/stretchr/testify/assert"
  16. "github.com/stretchr/testify/require"
  17. "gorm.io/gorm"
  18. "gorm.io/gorm/logger"
  19. )
  20. func setupDoubaoAssetProxyRouter(t *testing.T) *gin.Engine {
  21. t.Helper()
  22. oldMode := gin.Mode()
  23. gin.SetMode(gin.TestMode)
  24. t.Cleanup(func() {
  25. gin.SetMode(oldMode)
  26. })
  27. r := gin.New()
  28. r.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
  29. return r
  30. }
  31. func decodeDoubaoAssetErrorMessage(t *testing.T, body string) string {
  32. t.Helper()
  33. var payload struct {
  34. Error struct {
  35. Message string `json:"message"`
  36. Type string `json:"type"`
  37. } `json:"error"`
  38. }
  39. require.NoError(t, common.Unmarshal([]byte(body), &payload))
  40. return payload.Error.Message
  41. }
  42. func setupDoubaoAssetProxyDB(t *testing.T) *gorm.DB {
  43. t.Helper()
  44. oldDB := model.DB
  45. oldSQLitePath := common.SQLitePath
  46. oldMemoryCacheEnabled := common.MemoryCacheEnabled
  47. oldIsMasterNode := common.IsMasterNode
  48. oldUsingSQLite := common.UsingSQLite
  49. oldUsingMySQL := common.UsingMySQL
  50. oldUsingPostgreSQL := common.UsingPostgreSQL
  51. oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")
  52. common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared"
  53. common.MemoryCacheEnabled = false
  54. common.IsMasterNode = false
  55. common.UsingSQLite = false
  56. common.UsingMySQL = false
  57. common.UsingPostgreSQL = false
  58. require.NoError(t, os.Setenv("SQL_DSN", "local"))
  59. require.NoError(t, model.InitDB())
  60. model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
  61. db := model.DB
  62. sqlDB, err := db.DB()
  63. require.NoError(t, err)
  64. sqlDB.SetMaxOpenConns(1)
  65. require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}, &model.UserAssetGroup{}))
  66. t.Cleanup(func() {
  67. _ = sqlDB.Close()
  68. model.DB = oldDB
  69. common.SQLitePath = oldSQLitePath
  70. common.MemoryCacheEnabled = oldMemoryCacheEnabled
  71. common.IsMasterNode = oldIsMasterNode
  72. common.UsingSQLite = oldUsingSQLite
  73. common.UsingMySQL = oldUsingMySQL
  74. common.UsingPostgreSQL = oldUsingPostgreSQL
  75. if hadSQLDSN {
  76. _ = os.Setenv("SQL_DSN", oldSQLDSN)
  77. } else {
  78. _ = os.Unsetenv("SQL_DSN")
  79. }
  80. })
  81. return db
  82. }
  83. func createDoubaoAssetProxyChannel(t *testing.T, db *gorm.DB, id int, channelType int, group string, key string, status int) {
  84. t.Helper()
  85. priority := int64(id)
  86. weight := uint(10)
  87. autoBan := 1
  88. require.NoError(t, db.Create(&model.Channel{
  89. Id: id,
  90. Type: channelType,
  91. Key: key,
  92. Status: status,
  93. Name: fmt.Sprintf("channel-%d", id),
  94. Group: group,
  95. Models: "seedance-2",
  96. Priority: &priority,
  97. Weight: &weight,
  98. AutoBan: &autoBan,
  99. CreatedTime: int64(id),
  100. }).Error)
  101. }
  102. func TestDoubaoAssetProxyMissingActionReturns400(t *testing.T) {
  103. router := setupDoubaoAssetProxyRouter(t)
  104. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Version=2024-01-01", nil)
  105. w := httptest.NewRecorder()
  106. router.ServeHTTP(w, req)
  107. assert.Equal(t, http.StatusBadRequest, w.Code)
  108. assert.Equal(t, "Action query parameter is required", decodeDoubaoAssetErrorMessage(t, w.Body.String()))
  109. }
  110. func TestDoubaoAssetProxyUnknownActionReturns400(t *testing.T) {
  111. router := setupDoubaoAssetProxyRouter(t)
  112. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=CreateRealPersonAuthSession", nil)
  113. w := httptest.NewRecorder()
  114. router.ServeHTTP(w, req)
  115. assert.Equal(t, http.StatusBadRequest, w.Code)
  116. assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "unsupported asset Action")
  117. }
  118. func TestDoubaoAssetProxyInvalidJSONReturns400(t *testing.T) {
  119. router := setupDoubaoAssetProxyRouter(t)
  120. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"PageNumber":`))
  121. w := httptest.NewRecorder()
  122. router.ServeHTTP(w, req)
  123. assert.Equal(t, http.StatusBadRequest, w.Code)
  124. assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "invalid JSON body")
  125. }
  126. func TestDoubaoAssetProxyRejectsChinaMobileAssetGroupActions(t *testing.T) {
  127. db := setupDoubaoAssetProxyDB(t)
  128. fake := &fakeDoubaoAssetAdapter{}
  129. t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
  130. createDoubaoAssetProxyChannel(t, db, 61, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
  131. router := gin.New()
  132. router.Use(func(c *gin.Context) {
  133. c.Set("id", 10)
  134. common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
  135. c.Next()
  136. })
  137. router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
  138. for _, action := range []string{"CreateAssetGroup", "ListAssetGroups", "GetAssetGroup", "UpdateAssetGroup", "DeleteAssetGroup"} {
  139. t.Run(action, func(t *testing.T) {
  140. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, strings.NewReader(`{"Id":"group-1","GroupType":"AIGC","Name":"g"}`))
  141. w := httptest.NewRecorder()
  142. req.Header.Set("Content-Type", "application/json")
  143. router.ServeHTTP(w, req)
  144. assert.Equal(t, http.StatusForbidden, w.Code)
  145. assert.Equal(t, 0, fake.calls)
  146. })
  147. }
  148. }
  149. func TestDoubaoAssetProxyDoesNotBlockNonChinaMobileAssetGroupAction(t *testing.T) {
  150. db := setupDoubaoAssetProxyDB(t)
  151. fake := &fakeDoubaoAssetAdapter{}
  152. t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeDoubaoVideoCompatibleAiping, fake))
  153. createDoubaoAssetProxyChannel(t, db, 63, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "video-generation-key", common.ChannelStatusEnabled)
  154. router := newDoubaoAssetProxyTestRouter()
  155. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssetGroups", strings.NewReader(`{"PageNumber":1,"PageSize":10}`))
  156. w := httptest.NewRecorder()
  157. router.ServeHTTP(w, req)
  158. assert.Equal(t, http.StatusOK, w.Code)
  159. assert.Equal(t, service.AssetOperationAssetGroupList, fake.operation)
  160. }
  161. func TestDoubaoAssetProxyAutoGroupContinuesAfterMissingGroup(t *testing.T) {
  162. db := setupDoubaoAssetProxyDB(t)
  163. fake := &fakeDoubaoAssetAdapter{}
  164. t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
  165. createDoubaoAssetProxyChannel(t, db, 62, constant.ChannelTypeChinaMobileSeedance, "vip", "video-generation-key", common.ChannelStatusEnabled)
  166. router := gin.New()
  167. router.Use(func(c *gin.Context) {
  168. c.Set("id", 10)
  169. common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
  170. common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
  171. c.Next()
  172. })
  173. router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
  174. oldAutoGroups := doubaoAssetAutoGroups
  175. doubaoAssetAutoGroups = func(string) []string {
  176. return []string{"default", "vip"}
  177. }
  178. t.Cleanup(func() {
  179. doubaoAssetAutoGroups = oldAutoGroups
  180. })
  181. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
  182. w := httptest.NewRecorder()
  183. req.Header.Set("Content-Type", "application/json")
  184. router.ServeHTTP(w, req)
  185. assert.Equal(t, http.StatusOK, w.Code)
  186. assert.Equal(t, service.AssetOperationAssetList, fake.operation)
  187. }
  188. func TestDoubaoAssetProxyForwardsAdapterResponseHeaders(t *testing.T) {
  189. db := setupDoubaoAssetProxyDB(t)
  190. fake := &fakeDoubaoAssetAdapter{responseHeader: http.Header{"X-Upstream-Request-Id": []string{"upstream-1"}, "Content-Length": []string{"999"}}}
  191. t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
  192. createDoubaoAssetProxyChannel(t, db, 62, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
  193. router := gin.New()
  194. router.Use(func(c *gin.Context) {
  195. c.Set("id", 10)
  196. common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
  197. c.Next()
  198. })
  199. router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
  200. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
  201. w := httptest.NewRecorder()
  202. req.Header.Set("Content-Type", "application/json")
  203. router.ServeHTTP(w, req)
  204. assert.Equal(t, http.StatusOK, w.Code)
  205. assert.Equal(t, "upstream-1", w.Header().Get("X-Upstream-Request-Id"))
  206. assert.NotEqual(t, "999", w.Header().Get("Content-Length"))
  207. }
  208. func TestDoubaoAssetProxyScopesChinaMobileCreateAndListRequests(t *testing.T) {
  209. db := setupDoubaoAssetProxyDB(t)
  210. fake := &fakeDoubaoAssetAdapter{}
  211. t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
  212. createDoubaoAssetProxyChannel(t, db, 71, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
  213. require.NoError(t, model.CreateUserAssetGroup(10, 71, "owned"))
  214. router := newDoubaoAssetProxyTestRouter()
  215. for _, tc := range []struct {
  216. action string
  217. body string
  218. }{
  219. {"CreateAsset", `{"GroupId":"forged","Name":"n","URL":"https://example.com/a.png","AssetType":"Image"}`},
  220. {"ListAssets", `{"Filter":{"GroupType":"AIGC","GroupIds":["forged"]},"PageNumber":1,"PageSize":10}`},
  221. } {
  222. t.Run(tc.action, func(t *testing.T) {
  223. fake.requests = nil
  224. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+tc.action, strings.NewReader(tc.body))
  225. w := httptest.NewRecorder()
  226. router.ServeHTTP(w, req)
  227. require.Equal(t, http.StatusOK, w.Code)
  228. require.Len(t, fake.requests, 1)
  229. if tc.action == "CreateAsset" {
  230. assert.Equal(t, "owned", fake.requests[0].Body["GroupId"])
  231. } else {
  232. assert.Equal(t, []string{"owned"}, fake.requests[0].Body["Filter"].(map[string]any)["GroupIds"])
  233. }
  234. })
  235. }
  236. }
  237. func TestDoubaoAssetProxyRejectsChinaMobileListAssetsWithoutFilter(t *testing.T) {
  238. db := setupDoubaoAssetProxyDB(t)
  239. createDoubaoAssetProxyChannel(t, db, 73, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
  240. require.NoError(t, model.CreateUserAssetGroup(10, 73, "owned"))
  241. router := newDoubaoAssetProxyTestRouter()
  242. for _, body := range []string{`{"PageNumber":1,"PageSize":10}`, `{"Filter":null,"PageNumber":1,"PageSize":10}`} {
  243. w := httptest.NewRecorder()
  244. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(body))
  245. router.ServeHTTP(w, req)
  246. assert.Equal(t, http.StatusBadRequest, w.Code)
  247. assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "Filter.GroupType is required")
  248. }
  249. }
  250. func TestDoubaoAssetProxyHidesChinaMobileAssetsOwnedByAnotherUser(t *testing.T) {
  251. db := setupDoubaoAssetProxyDB(t)
  252. fake := &fakeDoubaoAssetAdapter{getAssetGroupID: "other"}
  253. t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
  254. createDoubaoAssetProxyChannel(t, db, 72, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
  255. require.NoError(t, model.CreateUserAssetGroup(10, 72, "owned"))
  256. router := newDoubaoAssetProxyTestRouter()
  257. for _, action := range []string{"GetAsset", "UpdateAsset", "DeleteAsset"} {
  258. t.Run(action, func(t *testing.T) {
  259. fake.requests = nil
  260. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, strings.NewReader(`{"Id":"asset-other","Name":"n"}`))
  261. w := httptest.NewRecorder()
  262. router.ServeHTTP(w, req)
  263. assert.Equal(t, http.StatusNotFound, w.Code)
  264. require.Len(t, fake.requests, 1)
  265. assert.Equal(t, service.AssetOperationAssetGet, fake.requests[0].Action.Operation)
  266. })
  267. }
  268. }
  269. func newDoubaoAssetProxyTestRouter() *gin.Engine {
  270. router := gin.New()
  271. router.Use(func(c *gin.Context) {
  272. c.Set("id", 10)
  273. common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
  274. c.Next()
  275. })
  276. router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
  277. return router
  278. }
  279. func TestDoubaoAssetProxyAutoGroupDoesNotBypassInvalidBinding(t *testing.T) {
  280. db := setupDoubaoAssetProxyDB(t)
  281. createDoubaoAssetProxyChannel(t, db, 1, constant.ChannelTypeChinaMobileSeedance, "other", "bad-key", common.ChannelStatusEnabled)
  282. createDoubaoAssetProxyChannel(t, db, 2, constant.ChannelTypeChinaMobileSeedance, "vip", "video-generation-key", common.ChannelStatusEnabled)
  283. require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeChinaMobileSeedance, "default", 1))
  284. router := gin.New()
  285. router.Use(func(c *gin.Context) {
  286. c.Set("id", 10)
  287. common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
  288. common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
  289. c.Next()
  290. })
  291. router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
  292. oldAutoGroups := doubaoAssetAutoGroups
  293. doubaoAssetAutoGroups = func(string) []string {
  294. return []string{"default", "vip"}
  295. }
  296. t.Cleanup(func() {
  297. doubaoAssetAutoGroups = oldAutoGroups
  298. })
  299. req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
  300. w := httptest.NewRecorder()
  301. req.Header.Set("Content-Type", "application/json")
  302. router.ServeHTTP(w, req)
  303. assert.Equal(t, http.StatusBadGateway, w.Code)
  304. assert.Contains(t, w.Body.String(), service.AssetErrorBindingInvalid)
  305. }
  306. func TestEffectiveDoubaoAssetGroupUsesUsingGroupBeforeBlankTokenGroup(t *testing.T) {
  307. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  308. common.SetContextKey(c, constant.ContextKeyTokenGroup, "")
  309. common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
  310. assert.Equal(t, "default", effectiveDoubaoAssetGroup(c))
  311. }
  312. func TestConcreteDoubaoAssetGroupsForAutoUsesUserAutoGroups(t *testing.T) {
  313. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  314. common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
  315. common.SetContextKey(c, constant.ContextKeyTokenGroup, "")
  316. groups := concreteDoubaoAssetGroupsForRequest(c, func(string) []string {
  317. return []string{"default", "vip"}
  318. })
  319. assert.Equal(t, []string{"default", "vip"}, groups)
  320. }
  321. type fakeDoubaoAssetAdapter struct {
  322. operation service.AssetOperation
  323. responseHeader http.Header
  324. requests []service.AssetRequest
  325. getAssetGroupID string
  326. calls int
  327. }
  328. func (a *fakeDoubaoAssetAdapter) Name() string {
  329. return "fake_asset"
  330. }
  331. func (a *fakeDoubaoAssetAdapter) Supports(operation service.AssetOperation) bool {
  332. return true
  333. }
  334. func (a *fakeDoubaoAssetAdapter) DoAssetRequest(_ context.Context, _ *model.Channel, req service.AssetRequest) (*service.AssetUpstreamResponse, *service.AssetError) {
  335. a.operation = req.Action.Operation
  336. a.requests = append(a.requests, req)
  337. a.calls++
  338. if req.Action.Operation == service.AssetOperationAssetGroupCreate {
  339. body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{"GroupId": "generated-group"})
  340. if err != nil {
  341. return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
  342. }
  343. return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
  344. }
  345. if req.Action.Operation == service.AssetOperationAssetGet {
  346. body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{"Id": "asset-1", "GroupId": a.getAssetGroupID})
  347. if err != nil {
  348. return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
  349. }
  350. return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
  351. }
  352. body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{
  353. "Items": []map[string]any{
  354. {"Id": "group-1", "Name": "g", "GroupType": "AIGC"},
  355. },
  356. "TotalCount": 1,
  357. })
  358. if err != nil {
  359. return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
  360. }
  361. return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
  362. }