feat: publish v1.0.0-rc.36-custom.1

controller/group_availability_test.go

@@ -0,0 +1,85 @@
+package controller
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"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/ratio_setting"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSelectableGroupsRequireEnabledChannel(t *testing.T) {
+	originalUsableGroups := setting.UserUsableGroups2JSONString()
+	originalRatios := ratio_setting.GroupRatio2JSONString()
+	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto":"Auto","active":"Active","default":"Default","paused":"Paused","no-channel":"No channel"}`))
+	require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"active":0.2,"default":1,"paused":0.3,"no-channel":0.4}`))
+	t.Cleanup(func() {
+		require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
+		require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
+		model.InvalidatePricingCache()
+	})
+
+	db := setupModelListControllerTestDB(t)
+	require.NoError(t, db.Create(&model.User{
+		Id:       7001,
+		Username: "group-availability-user",
+		Password: "password",
+		Group:    "default",
+		Status:   common.UserStatusEnabled,
+	}).Error)
+	require.NoError(t, db.Create(&[]model.Channel{
+		{Id: 11, Type: constant.ChannelTypeOpenAI, Key: "active-key", Name: "active", Status: common.ChannelStatusEnabled},
+		{Id: 12, Type: constant.ChannelTypeOpenAI, Key: "default-key", Name: "default", Status: common.ChannelStatusManuallyDisabled},
+		{Id: 13, Type: constant.ChannelTypeOpenAI, Key: "paused-key", Name: "paused", Status: common.ChannelStatusAutoDisabled},
+	}).Error)
+	require.NoError(t, db.Create(&[]model.Ability{
+		{Group: "active", Model: "gpt-4o", ChannelId: 11, Enabled: true},
+		{Group: "default", Model: "gpt-4o", ChannelId: 12, Enabled: true},
+		{Group: "paused", Model: "gpt-4o", ChannelId: 13, Enabled: true},
+	}).Error)
+	model.InvalidatePricingCache()
+
+	userRecorder := httptest.NewRecorder()
+	userContext, _ := gin.CreateTestContext(userRecorder)
+	userContext.Request = httptest.NewRequest(http.MethodGet, "/api/user/groups", nil)
+	userContext.Set("id", 7001)
+	GetUserGroups(userContext)
+
+	var userResponse struct {
+		Success bool                              `json:"success"`
+		Data    map[string]map[string]interface{} `json:"data"`
+	}
+	require.NoError(t, common.Unmarshal(userRecorder.Body.Bytes(), &userResponse))
+	require.True(t, userResponse.Success)
+	assert.ElementsMatch(t, []string{"active", "auto"}, mapKeys(userResponse.Data))
+
+	pricingRecorder := httptest.NewRecorder()
+	pricingContext, _ := gin.CreateTestContext(pricingRecorder)
+	pricingContext.Request = httptest.NewRequest(http.MethodGet, "/api/pricing", nil)
+	GetPricing(pricingContext)
+
+	var pricingResponse struct {
+		Success     bool               `json:"success"`
+		GroupRatio  map[string]float64 `json:"group_ratio"`
+		UsableGroup map[string]string  `json:"usable_group"`
+	}
+	require.NoError(t, common.Unmarshal(pricingRecorder.Body.Bytes(), &pricingResponse))
+	require.True(t, pricingResponse.Success)
+	assert.Equal(t, map[string]float64{"active": 0.2}, pricingResponse.GroupRatio)
+	assert.Equal(t, map[string]string{"active": "Active", "auto": "Auto"}, pricingResponse.UsableGroup)
+}
+
+func mapKeys[V any](values map[string]V) []string {
+	keys := make([]string, 0, len(values))
+	for key := range values {
+		keys = append(keys, key)
+	}
+	return keys
+}