73 lines
2.7 KiB
Go
73 lines
2.7 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/QuantumNous/new-api/common"
|
||
|
|
"github.com/QuantumNous/new-api/constant"
|
||
|
|
"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 configureRequestAutoGroupsTest(t *testing.T) {
|
||
|
|
t.Helper()
|
||
|
|
originalMax := setting.GetMaxTokenAutoGroups()
|
||
|
|
originalAutoGroups := setting.AutoGroups2JsonString()
|
||
|
|
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
||
|
|
originalRatios := ratio_setting.GroupRatio2JSONString()
|
||
|
|
require.NoError(t, setting.UpdateMaxTokenAutoGroups("2"))
|
||
|
|
require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["vip","default","svip"]`))
|
||
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","svip":"SVIP"}`))
|
||
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`))
|
||
|
|
t.Cleanup(func() {
|
||
|
|
require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax)))
|
||
|
|
require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups))
|
||
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
||
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func newRequestAutoGroupsContext() *gin.Context {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||
|
|
return ctx
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGetRequestAutoGroupsInheritedListIsNotLimited(t *testing.T) {
|
||
|
|
configureRequestAutoGroupsTest(t)
|
||
|
|
ctx := newRequestAutoGroupsContext()
|
||
|
|
|
||
|
|
groups := GetRequestAutoGroups(ctx, "default")
|
||
|
|
|
||
|
|
assert.Equal(t, []string{"vip", "default", "svip"}, groups)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGetRequestAutoGroupsFiltersBeforeApplyingCurrentLimit(t *testing.T) {
|
||
|
|
configureRequestAutoGroupsTest(t)
|
||
|
|
ctx := newRequestAutoGroupsContext()
|
||
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"revoked", "vip", "default", "svip"})
|
||
|
|
require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`))
|
||
|
|
|
||
|
|
groups := GetRequestAutoGroups(ctx, "default")
|
||
|
|
|
||
|
|
assert.Equal(t, []string{"vip", "default"}, groups)
|
||
|
|
require.NoError(t, setting.UpdateMaxTokenAutoGroups("1"))
|
||
|
|
assert.Equal(t, []string{"vip"}, GetRequestAutoGroups(ctx, "default"))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGetRequestAutoGroupsDoesNotFallBackAfterPermissionChange(t *testing.T) {
|
||
|
|
configureRequestAutoGroupsTest(t)
|
||
|
|
ctx := newRequestAutoGroupsContext()
|
||
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
|
||
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
|
||
|
|
|
||
|
|
groups := GetRequestAutoGroups(ctx, "default")
|
||
|
|
|
||
|
|
assert.Empty(t, groups)
|
||
|
|
}
|