// 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()) }