controller/token_auto_groups_test.go
Rama main · réplica de lectura del repositorio original.
package controller
import (
"fmt"
"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 configureTokenAutoGroupsTest(t *testing.T, maxCount string, autoGroups string) {
t.Helper()
originalMax := setting.GetMaxTokenAutoGroups()
originalAutoGroups := setting.AutoGroups2JsonString()
originalUsableGroups := setting.UserUsableGroups2JSONString()
originalRatios := ratio_setting.GroupRatio2JSONString()
require.NoError(t, setting.UpdateMaxTokenAutoGroups(maxCount))
require.NoError(t, setting.UpdateAutoGroupsByJsonString(autoGroups))
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`))
t.Cleanup(func() {
require.NoError(t, setting.UpdateMaxTokenAutoGroups(stringInt(originalMax)))
require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups))
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
})
}
func stringInt(value int) string {
return fmt.Sprintf("%d", value)
}
func setupTokenAutoGroupsControllerTest(t *testing.T) *model.User {
t.Helper()
db := setupTokenControllerTestDB(t)
require.NoError(t, db.AutoMigrate(&model.User{}))
user := &model.User{
Id: 101,
Username: "token-auto-user",
Password: "password",
Group: "default",
Status: common.UserStatusEnabled,
}
require.NoError(t, db.Create(user).Error)
return user
}
func baseAutoTokenRequest(name string) map[string]any {
return map[string]any{
"name": name,
"expired_time": -1,
"remain_quota": 0,
"unlimited_quota": true,
"group": "auto",
"cross_group_retry": true,
}
}
func newTokenAutoGroupsAuthenticatedContext(t *testing.T, method string, target string, body any, userID int) (*gin.Context, *httptest.ResponseRecorder) {
t.Helper()
ctx, recorder := newAuthenticatedContext(t, method, target, body, userID)
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
return ctx, recorder
}
func TestAddTokenEmptyAutoGroupsInheritGlobalAuto(t *testing.T) {
tests := []struct {
name string
includeField bool
value any
}{
{name: "omitted"},
{name: "null", includeField: true, value: nil},
{name: "empty array", includeField: true, value: []string{}},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
configureTokenAutoGroupsTest(t, "5", `["default","vip"]`)
user := setupTokenAutoGroupsControllerTest(t)
request := baseAutoTokenRequest("create-" + test.name)
if test.includeField {
request["auto_groups"] = test.value
}
ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id)
AddToken(ctx)
response := decodeAPIResponse(t, recorder)
require.True(t, response.Success, response.Message)
var token model.Token
require.NoError(t, model.DB.Where("name = ?", request["name"]).First(&token).Error)
assert.Empty(t, token.AutoGroups)
assert.True(t, token.CrossGroupRetry)
payload, err := common.Marshal(buildMaskedTokenResponse(&token))
require.NoError(t, err)
var responseData map[string]any
require.NoError(t, common.Unmarshal(payload, &responseData))
assert.Nil(t, responseData["auto_groups"])
})
}
}
func TestAddTokenPersistsOrderedAutoGroupsSnapshot(t *testing.T) {
configureTokenAutoGroupsTest(t, "5", `["default","vip"]`)
user := setupTokenAutoGroupsControllerTest(t)
request := baseAutoTokenRequest("ordered-snapshot")
request["auto_groups"] = []string{"vip", "default"}
ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id)
AddToken(ctx)
require.True(t, decodeAPIResponse(t, recorder).Success)
var token model.Token
require.NoError(t, model.DB.Where("name = ?", "ordered-snapshot").First(&token).Error)
assert.JSONEq(t, `["vip","default"]`, token.AutoGroups)
getCtx, getRecorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/"+stringInt(token.Id), nil, user.Id)
getCtx.Params = append(getCtx.Params, gin.Param{Key: "id", Value: stringInt(token.Id)})
GetToken(getCtx)
getResponse := decodeAPIResponse(t, getRecorder)
require.True(t, getResponse.Success)
var data struct {
AutoGroups []string `json:"auto_groups"`
}
require.NoError(t, common.Unmarshal(getResponse.Data, &data))
assert.Equal(t, []string{"vip", "default"}, data.AutoGroups)
}
func TestUpdateTokenAutoGroupsTriStateAndNonAutoCleanup(t *testing.T) {
tests := []struct {
name string
includeField bool
value any
group string
expectedAutoGroups string
expectedRetry bool
}{
{name: "omitted preserves", group: "auto", expectedAutoGroups: `["vip","default"]`, expectedRetry: true},
{name: "null inherits", includeField: true, value: nil, group: "auto", expectedRetry: true},
{name: "empty inherits", includeField: true, value: []string{}, group: "auto", expectedRetry: true},
{name: "non auto clears and disables retry", includeField: true, value: []string{"vip"}, group: "default"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
configureTokenAutoGroupsTest(t, "5", `["default","vip"]`)
user := setupTokenAutoGroupsControllerTest(t)
token := seedToken(t, model.DB, user.Id, "update-auto", "update-auto-key")
token.Group = "auto"
token.CrossGroupRetry = true
require.NoError(t, token.SetAutoGroups([]string{"vip", "default"}))
require.NoError(t, model.DB.Save(token).Error)
request := baseAutoTokenRequest("updated-auto")
request["id"] = token.Id
request["status"] = common.TokenStatusEnabled
request["group"] = test.group
if test.includeField {
request["auto_groups"] = test.value
}
ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPut, "/api/token/", request, user.Id)
UpdateToken(ctx)
response := decodeAPIResponse(t, recorder)
require.True(t, response.Success, response.Message)
var updated model.Token
require.NoError(t, model.DB.First(&updated, token.Id).Error)
if test.expectedAutoGroups == "" {
assert.Empty(t, updated.AutoGroups)
} else {
assert.JSONEq(t, test.expectedAutoGroups, updated.AutoGroups)
}
assert.Equal(t, test.expectedRetry, updated.CrossGroupRetry)
})
}
}
func TestAddTokenRejectsInvalidAutoGroups(t *testing.T) {
tests := []struct {
name string
maxCount string
groups []string
}{
{name: "over limit", maxCount: "1", groups: []string{"default", "vip"}},
{name: "duplicate", maxCount: "5", groups: []string{"default", "default"}},
{name: "auto pseudo group", maxCount: "5", groups: []string{"auto"}},
{name: "unavailable", maxCount: "5", groups: []string{"missing"}},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
configureTokenAutoGroupsTest(t, test.maxCount, `["default","vip"]`)
user := setupTokenAutoGroupsControllerTest(t)
request := baseAutoTokenRequest("invalid-" + test.name)
request["auto_groups"] = test.groups
ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id)
AddToken(ctx)
response := decodeAPIResponse(t, recorder)
assert.False(t, response.Success)
var count int64
require.NoError(t, model.DB.Model(&model.Token{}).Count(&count).Error)
assert.Zero(t, count)
})
}
}
func TestGetTokenAutoGroupsReturnsFullFilteredGlobalOrderAndLimit(t *testing.T) {
configureTokenAutoGroupsTest(t, "1", `["vip","missing","default"]`)
user := setupTokenAutoGroupsControllerTest(t)
ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/auto-groups", nil, user.Id)
GetTokenAutoGroups(ctx)
response := decodeAPIResponse(t, recorder)
require.True(t, response.Success, response.Message)
var data struct {
Groups []string `json:"groups"`
MaxCount int `json:"max_count"`
}
require.NoError(t, common.Unmarshal(response.Data, &data))
assert.Equal(t, []string{"vip", "default"}, data.Groups)
assert.Equal(t, 1, data.MaxCount)
}