118 lines
4.2 KiB
Go
118 lines
4.2 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestTokenAstraCacheUpgradePreservesLiveQuota(t *testing.T) {
|
|
useUserCacheMiniRedis(t)
|
|
ctx := context.Background()
|
|
token := Token{Id: 43, UserId: 7, Key: "legacy-astra-cache", AstraAutoSwitch: true, RemainQuota: 100}
|
|
require.NoError(t, cacheSetTokenForTest(token))
|
|
key := getTokenCacheKey(token.Key)
|
|
require.NoError(t, common.RDB.HDel(ctx, key, "AstraAutoSwitch").Err())
|
|
require.NoError(t, common.RDB.HSet(ctx, key, "RemainQuota", 70).Err())
|
|
_, err := cacheGetTokenByKey(token.Key)
|
|
require.Error(t, err, "legacy cache must fetch the migrated value from the database")
|
|
require.NoError(t, cacheSetTokenForTest(token))
|
|
cached, err := cacheGetTokenByKey(token.Key)
|
|
require.NoError(t, err)
|
|
require.True(t, cached.AstraAutoSwitch)
|
|
require.Equal(t, 70, cached.RemainQuota, "metadata hydration must not overwrite concurrent quota consumption")
|
|
}
|
|
|
|
func TestTokenAutoGroupsRoundTripThroughRedisHashCache(t *testing.T) {
|
|
useUserCacheMiniRedis(t)
|
|
token := Token{
|
|
Id: 42,
|
|
UserId: 7,
|
|
Key: "token-auto-groups-cache-key",
|
|
Name: "auto-cache",
|
|
Group: "auto",
|
|
AutoGroups: `["vip","default"]`,
|
|
}
|
|
|
|
require.NoError(t, cacheSetTokenForTest(token))
|
|
cached, err := cacheGetTokenByKey(token.Key)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, token.AutoGroups, cached.AutoGroups)
|
|
groups, err := cached.GetAutoGroups()
|
|
require.NoError(t, err)
|
|
assert.Equal(t, []string{"vip", "default"}, groups)
|
|
}
|
|
|
|
func TestTokenAllowUserQuotaQueryRoundTripsThroughRedisHashCache(t *testing.T) {
|
|
server := useUserCacheMiniRedis(t)
|
|
token := Token{Id: 43, UserId: 7, Key: "token-user-quota-cache-key", Name: "user-quota-cache", AllowUserQuotaQuery: true}
|
|
|
|
initialized, err := cacheInitToken(token)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, initialized)
|
|
cached, err := cacheGetTokenByKey(token.Key)
|
|
require.NoError(t, err)
|
|
assert.True(t, cached.AllowUserQuotaQuery)
|
|
assert.Positive(t, server.TTL(getTokenCacheKey(token.Key)))
|
|
}
|
|
|
|
func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) {
|
|
truncateTables(t)
|
|
useUserCacheMiniRedis(t)
|
|
token := Token{
|
|
UserId: 7,
|
|
Key: "token-auto-groups-update-cache-key",
|
|
Name: "auto-cache-update",
|
|
Status: common.TokenStatusEnabled,
|
|
ExpiredTime: -1,
|
|
UnlimitedQuota: true,
|
|
Group: "auto",
|
|
CrossGroupRetry: true,
|
|
AutoGroups: `["default","vip"]`,
|
|
AstraAutoSwitch: true,
|
|
}
|
|
require.NoError(t, token.Insert())
|
|
require.NoError(t, cacheSetTokenForTest(token))
|
|
|
|
preheated, err := cacheGetTokenByKey(token.Key)
|
|
require.NoError(t, err)
|
|
assert.JSONEq(t, `["default","vip"]`, preheated.AutoGroups)
|
|
|
|
require.NoError(t, token.SetAutoGroups([]string{"vip"}))
|
|
token.AstraAutoSwitch = false
|
|
require.NoError(t, token.Update())
|
|
// Update 是限制性变更：写库前删除缓存并设置 fence。缓存不再提供旧的
|
|
// 宽分组值，下一次读取必须看到收紧后的分组。
|
|
_, cacheErr := cacheGetTokenByKey(token.Key)
|
|
require.Error(t, cacheErr, "the pre-update cache entry must be invalidated")
|
|
reloaded, err := GetTokenByKey(token.Key, false)
|
|
require.NoError(t, err)
|
|
assert.JSONEq(t, `["vip"]`, reloaded.AutoGroups)
|
|
assert.False(t, reloaded.AstraAutoSwitch)
|
|
}
|
|
|
|
// cacheSetTokenForTest 以测试身份写入完整 token 缓存（含额度字段），
|
|
// 模拟“已水合”的缓存状态。
|
|
func cacheSetTokenForTest(token Token) error {
|
|
_, err := cacheInitToken(token)
|
|
return err
|
|
}
|
|
|
|
func TestTokenRoutingBlacklistRoundTripsThroughRedisHashCache(t *testing.T) {
|
|
server := useUserCacheMiniRedis(t)
|
|
token := Token{Id: 43, UserId: 7, Key: "token-blacklist-cache-key", Name: "blacklist-cache", RoutingPolicy: `{"version":1,"mode":"smart","excluded_groups":["blocked"]}`}
|
|
|
|
initialized, err := cacheInitToken(token)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, initialized)
|
|
cached, err := cacheGetTokenByKey(token.Key)
|
|
require.NoError(t, err)
|
|
policy, err := cached.GetRoutingPolicy()
|
|
require.NoError(t, err)
|
|
assert.Equal(t, []string{"blocked"}, policy.ExcludedGroups)
|
|
assert.Positive(t, server.TTL(getTokenCacheKey(token.Key)))
|
|
}
|