cf5160c9f2
Expose admin Redis snapshot for user and channel concurrency plus current user RPM usage.
117 lines
2.9 KiB
Go
117 lines
2.9 KiB
Go
package service
|
|
|
|
import (
|
|
"strconv"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestAcquireConcurrencyAllowsWhenLimitDisabled(t *testing.T) {
|
|
previousRedisEnabled := common.RedisEnabled
|
|
previousRDB := common.RDB
|
|
common.RedisEnabled = true
|
|
common.RDB = nil
|
|
t.Cleanup(func() {
|
|
common.RedisEnabled = previousRedisEnabled
|
|
common.RDB = previousRDB
|
|
})
|
|
|
|
allowed, release := AcquireUserConcurrency(1, 0)
|
|
require.True(t, allowed)
|
|
require.NotNil(t, release)
|
|
release()
|
|
|
|
allowed, release = AcquireChannelConcurrency(1, -1)
|
|
require.True(t, allowed)
|
|
require.NotNil(t, release)
|
|
release()
|
|
|
|
require.True(t, CheckUserRPM(1, 0))
|
|
}
|
|
|
|
func TestAcquireConcurrencyAllowsWhenRedisUnavailable(t *testing.T) {
|
|
for _, redisEnabled := range []bool{false, true} {
|
|
t.Run("redis_enabled_"+strconv.FormatBool(redisEnabled), func(t *testing.T) {
|
|
previousRedisEnabled := common.RedisEnabled
|
|
previousRDB := common.RDB
|
|
common.RedisEnabled = redisEnabled
|
|
common.RDB = nil
|
|
t.Cleanup(func() {
|
|
common.RedisEnabled = previousRedisEnabled
|
|
common.RDB = previousRDB
|
|
})
|
|
|
|
allowed, release := AcquireUserConcurrency(1, 1)
|
|
require.True(t, allowed)
|
|
require.NotNil(t, release)
|
|
release()
|
|
|
|
allowed, release = AcquireChannelConcurrency(1, 1)
|
|
require.True(t, allowed)
|
|
require.NotNil(t, release)
|
|
release()
|
|
|
|
require.True(t, CheckUserRPM(1, 1))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGetOpsConcurrencySnapshotReturnsEmptyWhenRedisDisabled(t *testing.T) {
|
|
previousRedisEnabled := common.RedisEnabled
|
|
previousRDB := common.RDB
|
|
common.RedisEnabled = false
|
|
common.RDB = nil
|
|
t.Cleanup(func() {
|
|
common.RedisEnabled = previousRedisEnabled
|
|
common.RDB = previousRDB
|
|
})
|
|
|
|
snapshot := GetOpsConcurrencySnapshot()
|
|
|
|
require.False(t, snapshot.RedisEnabled)
|
|
require.Empty(t, snapshot.Users)
|
|
require.Empty(t, snapshot.Channels)
|
|
require.NotZero(t, snapshot.GeneratedAt)
|
|
}
|
|
|
|
func TestParseOpsConcurrencyKeys(t *testing.T) {
|
|
userID, ok := parseUserConcurrencyKey("user:concurrency:123")
|
|
require.True(t, ok)
|
|
require.Equal(t, 123, userID)
|
|
|
|
channelID, ok := parseChannelConcurrencyKey("channel:concurrency:456")
|
|
require.True(t, ok)
|
|
require.Equal(t, 456, channelID)
|
|
|
|
rpmUserID, ok := parseUserRPMKey("user:rpm:789:202605201430", "202605201430")
|
|
require.True(t, ok)
|
|
require.Equal(t, 789, rpmUserID)
|
|
}
|
|
|
|
func TestParseOpsConcurrencyKeysRejectsInvalidKeys(t *testing.T) {
|
|
invalidUserKeys := []string{
|
|
"user:concurrency:",
|
|
"user:concurrency:0",
|
|
"user:concurrency:abc",
|
|
"channel:concurrency:1",
|
|
}
|
|
for _, key := range invalidUserKeys {
|
|
_, ok := parseUserConcurrencyKey(key)
|
|
require.False(t, ok)
|
|
}
|
|
|
|
invalidRPMKeys := []string{
|
|
"user:rpm:",
|
|
"user:rpm:0:202605201430",
|
|
"user:rpm:abc:202605201430",
|
|
"user:rpm:123:202605201429",
|
|
"user:rpm:123",
|
|
}
|
|
for _, key := range invalidRPMKeys {
|
|
_, ok := parseUserRPMKey(key, "202605201430")
|
|
require.False(t, ok)
|
|
}
|
|
}
|