feat: sync deployed v1.0.0-rc.38-custom.67

service/image_generation_group_test.go

@@ -0,0 +1,55 @@
+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 = "example-images", []string{"gpt-image-2"}
+	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","example-images":"Images"}`))
+	require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2,"example-images":1.25}`))
+	createChannelSelectAutoGroupsChannel(t, db, 8001, "default", "gpt-image-2")
+	createChannelSelectAutoGroupsChannel(t, db, 8002, "example-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: "example-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, "example-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"))
+}