41 lines
963 B
Go
41 lines
963 B
Go
// stm: #unit
|
|
package main
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestRateLimit(t *testing.T) {
|
|
//stm: @CMD_LIMITER_GET_IP_LIMITER_001, @CMD_LIMITER_GET_WALLET_LIMITER_001
|
|
limiter := NewLimiter(LimiterConfig{
|
|
TotalRate: time.Second,
|
|
TotalBurst: 20,
|
|
IPRate: time.Second,
|
|
IPBurst: 1,
|
|
WalletRate: time.Second,
|
|
WalletBurst: 1,
|
|
})
|
|
|
|
for i := 0; i < 20; i++ {
|
|
assert.True(t, limiter.Allow())
|
|
}
|
|
|
|
assert.False(t, limiter.Allow())
|
|
|
|
time.Sleep(time.Second)
|
|
assert.True(t, limiter.Allow())
|
|
|
|
assert.True(t, limiter.GetIPLimiter("127.0.0.1").Allow())
|
|
assert.False(t, limiter.GetIPLimiter("127.0.0.1").Allow())
|
|
time.Sleep(time.Second)
|
|
assert.True(t, limiter.GetIPLimiter("127.0.0.1").Allow())
|
|
|
|
assert.True(t, limiter.GetWalletLimiter("abc123").Allow())
|
|
assert.False(t, limiter.GetWalletLimiter("abc123").Allow())
|
|
time.Sleep(time.Second)
|
|
assert.True(t, limiter.GetWalletLimiter("abc123").Allow())
|
|
}
|