server
/
new-api
1
0
Bifurcation 0
Code Tickets 0 Demandes d'ajout 0 Versions 0 Wiki Activité
Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.
289 Révisions
2 Branches
352 MiB
 
 
 
Aborescence: 4533c24fa1
Branches Tags
${ item.name }
Créer la branche ${ searchTerm }
de '4533c24fa1'
${ noResults }
new-api/service/doubao_asset_channel_test.go

266 lignes
11 KiB
Brut Annotations Historique 

  1. package service

  2. 

  3. import (

  4. 	"fmt"

  5. 	"os"

  6. 	"strings"

  7. 	"testing"

  8. 

  9. 	"github.com/QuantumNous/new-api/common"

  10. 	"github.com/QuantumNous/new-api/constant"

  11. 	"github.com/QuantumNous/new-api/model"

  12. 	"github.com/stretchr/testify/assert"

  13. 	"github.com/stretchr/testify/require"

  14. 	"gorm.io/gorm"

  15. 	"gorm.io/gorm/logger"

  16. )

  17. 

  18. func setupDoubaoAssetChannelDB(t *testing.T) *gorm.DB {

  19. 	t.Helper()

  20. 

  21. 	oldDB := model.DB

  22. 	oldSQLitePath := common.SQLitePath

  23. 	oldMemoryCacheEnabled := common.MemoryCacheEnabled

  24. 	oldIsMasterNode := common.IsMasterNode

  25. 	oldUsingSQLite := common.UsingSQLite

  26. 	oldUsingMySQL := common.UsingMySQL

  27. 	oldUsingPostgreSQL := common.UsingPostgreSQL

  28. 	oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")

  29. 

  30. 	common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared"

  31. 	common.MemoryCacheEnabled = false

  32. 	common.IsMasterNode = false

  33. 	common.UsingSQLite = false

  34. 	common.UsingMySQL = false

  35. 	common.UsingPostgreSQL = false

  36. 	require.NoError(t, os.Setenv("SQL_DSN", "local"))

  37. 	require.NoError(t, model.InitDB())

  38. 

  39. 	model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})

  40. 	db := model.DB

  41. 

  42. 	sqlDB, err := db.DB()

  43. 	require.NoError(t, err)

  44. 	sqlDB.SetMaxOpenConns(1)

  45. 	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}))

  46. 

  47. 	t.Cleanup(func() {

  48. 		_ = sqlDB.Close()

  49. 		model.DB = oldDB

  50. 		common.SQLitePath = oldSQLitePath

  51. 		common.MemoryCacheEnabled = oldMemoryCacheEnabled

  52. 		common.IsMasterNode = oldIsMasterNode

  53. 		common.UsingSQLite = oldUsingSQLite

  54. 		common.UsingMySQL = oldUsingMySQL

  55. 		common.UsingPostgreSQL = oldUsingPostgreSQL

  56. 		if hadSQLDSN {

  57. 			_ = os.Setenv("SQL_DSN", oldSQLDSN)

  58. 		} else {

  59. 			_ = os.Unsetenv("SQL_DSN")

  60. 		}

  61. 	})

  62. 

  63. 	return db

  64. }

  65. 

  66. func createDoubaoAssetChannelForTest(t *testing.T, db *gorm.DB, id int, channelType int, group string, key string, status int) {

  67. 	t.Helper()

  68. 

  69. 	priority := int64(id)

  70. 	weight := uint(10)

  71. 	autoBan := 1

  72. 	channel := &model.Channel{

  73. 		Id:          id,

  74. 		Type:        channelType,

  75. 		Key:         key,

  76. 		Status:      status,

  77. 		Name:        fmt.Sprintf("channel-%d", id),

  78. 		Group:       group,

  79. 		Models:      "seedance-2",

  80. 		Priority:    &priority,

  81. 		Weight:      &weight,

  82. 		AutoBan:     &autoBan,

  83. 		CreatedTime: int64(id),

  84. 	}

  85. 	require.NoError(t, db.Create(channel).Error)

  86. }

  87. 

  88. func createDoubaoAssetAbilityForTest(t *testing.T, db *gorm.DB, group string, modelName string, channelId int, enabled bool) {

  89. 	t.Helper()

  90. 

  91. 	priority := int64(channelId)

  92. 	require.NoError(t, db.Create(&model.Ability{

  93. 		Group:     group,

  94. 		Model:     modelName,

  95. 		ChannelId: channelId,

  96. 		Enabled:   enabled,

  97. 		Priority:  &priority,

  98. 		Weight:    10,

  99. 	}).Error)

  100. }

  101. 

  102. func TestResolveDoubaoAssetChannel_PrefersBoundChannelWithKey(t *testing.T) {

  103. 	db := setupDoubaoAssetChannelDB(t)

  104. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "bound-key", common.ChannelStatusEnabled)

  105. 	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled)

  106. 	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1))

  107. 

  108. 	channel, err := ResolveDoubaoAssetChannel(10, "default")

  109. 

  110. 	require.NoError(t, err)

  111. 	require.NotNil(t, channel)

  112. 	assert.Equal(t, 1, channel.Id)

  113. 	assert.Equal(t, "bound-key", channel.Key)

  114. }

  115. 

  116. func TestResolveDoubaoAssetChannel_AutoMatchesAndPersistsWhenUnbound(t *testing.T) {

  117. 	db := setupDoubaoAssetChannelDB(t)

  118. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled)

  119. 

  120. 	channel, err := ResolveDoubaoAssetChannel(10, "default")

  121. 

  122. 	require.NoError(t, err)

  123. 	require.NotNil(t, channel)

  124. 	assert.Equal(t, 1, channel.Id)

  125. 

  126. 	binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default")

  127. 	require.NoError(t, err)

  128. 	require.NotNil(t, binding)

  129. 	assert.Equal(t, 1, binding.ChannelId)

  130. }

  131. 

  132. func TestResolveDoubaoAssetChannel_AutoMatchesTianyiYunAndPersistsByType(t *testing.T) {

  133. 	db := setupDoubaoAssetChannelDB(t)

  134. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled)

  135. 	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 1, true)

  136. 

  137. 	channel, err := ResolveDoubaoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0")

  138. 

  139. 	require.NoError(t, err)

  140. 	require.NotNil(t, channel)

  141. 	assert.Equal(t, 1, channel.Id)

  142. 	assert.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, channel.Type)

  143. 

  144. 	binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default")

  145. 	require.NoError(t, err)

  146. 	require.NotNil(t, binding)

  147. 	assert.Equal(t, 1, binding.ChannelId)

  148. }

  149. 

  150. func TestGetBoundDoubaoAssetChannelForModelFindsTianyiYunBindingForPublicModel(t *testing.T) {

  151. 	db := setupDoubaoAssetChannelDB(t)

  152. 	createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled)

  153. 	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true)

  154. 	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16))

  155. 

  156. 	channel, err := GetBoundDoubaoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0")

  157. 

  158. 	require.NoError(t, err)

  159. 	require.NotNil(t, channel)

  160. 	assert.Equal(t, 16, channel.Id)

  161. }

  162. 

  163. func TestResolveDoubaoAssetChannel_RebindsOpenAIMisbindToAiping(t *testing.T) {

  164. 	db := setupDoubaoAssetChannelDB(t)

  165. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled)

  166. 	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled)

  167. 	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1))

  168. 

  169. 	channel, err := ResolveDoubaoAssetChannel(10, "default")

  170. 

  171. 	require.NoError(t, err)

  172. 	require.NotNil(t, channel)

  173. 	assert.Equal(t, 2, channel.Id)

  174. 

  175. 	binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default")

  176. 	require.NoError(t, err)

  177. 	require.NotNil(t, binding)

  178. 	assert.Equal(t, 2, binding.ChannelId)

  179. }

  180. 

  181. func TestResolveDoubaoAssetChannel_RebindsOnlyCurrentGroupWhenBoundGroupMismatches(t *testing.T) {

  182. 	db := setupDoubaoAssetChannelDB(t)

  183. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", "vip-key", common.ChannelStatusEnabled)

  184. 	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled)

  185. 	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1))

  186. 	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", 1))

  187. 

  188. 	channel, err := ResolveDoubaoAssetChannel(10, "default")

  189. 

  190. 	require.NoError(t, err)

  191. 	require.NotNil(t, channel)

  192. 	assert.Equal(t, 2, channel.Id)

  193. 

  194. 	defaultBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default")

  195. 	require.NoError(t, err)

  196. 	require.NotNil(t, defaultBinding)

  197. 	assert.Equal(t, 2, defaultBinding.ChannelId)

  198. 

  199. 	vipBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip")

  200. 	require.NoError(t, err)

  201. 	require.NotNil(t, vipBinding)

  202. 	assert.Equal(t, 1, vipBinding.ChannelId)

  203. }

  204. 

  205. func TestResolveDoubaoAssetChannel_AutoMatchScansPastFirstPage(t *testing.T) {

  206. 	db := setupDoubaoAssetChannelDB(t)

  207. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled)

  208. 	for i := 2; i <= DoubaoAssetChannelPageSize+1; i++ {

  209. 		createDoubaoAssetChannelForTest(t, db, i, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", fmt.Sprintf("vip-key-%d", i), common.ChannelStatusEnabled)

  210. 	}

  211. 

  212. 	channel, err := ResolveDoubaoAssetChannel(10, "default")

  213. 

  214. 	require.NoError(t, err)

  215. 	require.NotNil(t, channel)

  216. 	assert.Equal(t, 1, channel.Id)

  217. }

  218. 

  219. func TestResolveDoubaoAssetChannel_ReturnsErrorWhenBindingWriteFails(t *testing.T) {

  220. 	db := setupDoubaoAssetChannelDB(t)

  221. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled)

  222. 

  223. 	const callbackName = "doubao_asset_test:fail_user_asset_channel_create"

  224. 	require.NoError(t, db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) {

  225. 		if _, ok := tx.Statement.Dest.(*model.UserAssetChannel); ok {

  226. 			tx.AddError(assert.AnError)

  227. 		}

  228. 	}))

  229. 

  230. 	channel, err := ResolveDoubaoAssetChannel(10, "default")

  231. 

  232. 	require.Error(t, err)

  233. 	assert.Nil(t, channel)

  234. }

  235. 

  236. func TestHasDoubaoVideoChannelForGroupModel(t *testing.T) {

  237. 	db := setupDoubaoAssetChannelDB(t)

  238. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled)

  239. 	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", "vip-key", common.ChannelStatusEnabled)

  240. 	createDoubaoAssetChannelForTest(t, db, 3, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled)

  241. 	createDoubaoAssetChannelForTest(t, db, 4, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled)

  242. 	createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 1, true)

  243. 	createDoubaoAssetAbilityForTest(t, db, "vip", "seedance-2", 2, true)

  244. 	createDoubaoAssetAbilityForTest(t, db, "default", "disabled-model", 1, false)

  245. 	createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 3, true)

  246. 	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 4, true)

  247. 

  248. 	assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "seedance-2"))

  249. 	assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "Doubao-Seedance-2.0"))

  250. 	assert.False(t, HasDoubaoVideoChannelForGroupModel("default", "disabled-model"))

  251. 	assert.False(t, HasDoubaoVideoChannelForGroupModel("missing", "seedance-2"))

  252. }

  253. 

  254. func TestAutoMatchDoubaoAssetChannel_IgnoresLegacyDoubaoVideoChannels(t *testing.T) {

  255. 	db := setupDoubaoAssetChannelDB(t)

  256. 	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideo, "default", "legacy-key", common.ChannelStatusEnabled)

  257. 	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled)

  258. 

  259. 	channel, err := AutoMatchDoubaoAssetChannel("default")

  260. 

  261. 	require.NoError(t, err)

  262. 	require.NotNil(t, channel)

  263. 	assert.Equal(t, 2, channel.Id)

  264. 	assert.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleAiping, channel.Type)

  265. }