f069f98b35
Add Redis-backed user and channel concurrency slots plus user RPM checks with fail-open behavior.
60 lines
1.4 KiB
Go
60 lines
1.4 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))
|
|
})
|
|
}
|
|
}
|