56 lines
3.0 KiB
Go
56 lines
3.0 KiB
Go
package service
|
|
|
|
import (
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
"github.com/QuantumNous/new-api/setting/model_setting"
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestDedicatedImageGenerationPoolSelection(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
cfg := model_setting.GetGlobalSettings()
|
|
oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
|
|
t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
|
|
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "gpt-images", []string{"gpt-image-2"}
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","gpt-images":"Images"}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2,"gpt-images":1.25}`))
|
|
createChannelSelectAutoGroupsChannel(t, db, 8001, "default", "gpt-image-2")
|
|
createChannelSelectAutoGroupsChannel(t, db, 8002, "gpt-images", "gpt-image-2")
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
for _, mode := range []string{"smart", "inherit", "custom"} {
|
|
t.Run(mode, func(t *testing.T) {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("POST", "/v1/images/generations", nil)
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}})
|
|
entries := ResolveRequestPool(c, "gpt-image-2")
|
|
require.Equal(t, []model.TokenGroupEntry{{Group: "gpt-images", Priority: 0}}, entries)
|
|
channel, group, err := CacheGetRandomSatisfiedChannel(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: "gpt-image-2", Retry: common.GetPointer(0)})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, channel)
|
|
require.Equal(t, 8002, channel.Id)
|
|
require.Equal(t, "gpt-images", group)
|
|
})
|
|
}
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("POST", "/v1/images/generations", nil)
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
|
|
defaults := []model.TokenGroupEntry{{Group: "default", Priority: 0}}
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: defaults})
|
|
require.Equal(t, defaults, ResolveRequestPool(c, "text-model"))
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "user"})
|
|
require.Equal(t, defaults, ResolveRequestPool(c, "gpt-image-2"))
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: defaults})
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
|
|
require.Equal(t, defaults, ResolveRequestPool(c, "gpt-image-2"))
|
|
}
|