39 lines
872 B
Go
39 lines
872 B
Go
package main
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestRateLimit(t *testing.T) {
|
|
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())
|
|
}
|