mirror of
https://gh-proxy.org/https://github.com/klinsmaya/new-api.git
synced 2026-11-08 01:17:09 +08:00
CaIon
7de10f2c60
Add a per-token "Allow querying user quota" switch. When it is on, GET /api/usage/token/ also returns user_quota with the account's remaining and used quota in quota units. New tokens start with the switch on: the console form, POST /api/token/ without the field, and the default token created at registration. Tokens created before this change stay off. PUT /api/token/ keeps the stored value when the field is absent, and the token audit lists the field when it changes. The key usage and key log queries now check the key's IP whitelist the same way API requests do, and answer 403 from an address outside it. Expired and exhausted keys can still use both queries. New API channels now query their balance from the upstream's /api/usage/token/ with the channel key. When the upstream also returns the owner's quota, the smaller of the key's and the owner's remaining quota is the balance, and an unlimited key uses the owner's quota alone. Without it, a limited key reports its own remaining quota, and an unlimited key asks for the switch to be turned on upstream. Quota converts to USD with the upstream's quota_per_unit from /api/status. Refs #3122 Refs #6828
82 lines
2.6 KiB
Go
82 lines
2.6 KiB
Go
package model
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
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"]`,
|
|
}
|
|
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"}))
|
|
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)
|
|
}
|
|
|
|
// cacheSetTokenForTest 以测试身份写入完整 token 缓存（含额度字段），
|
|
// 模拟“已水合”的缓存状态。
|
|
func cacheSetTokenForTest(token Token) error {
|
|
_, err := cacheInitToken(token)
|
|
return err
|
|
}
|