Ilia Denisov
50 lines
1011 B
Go
50 lines
1011 B
Go
package ratelimit
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestInMemoryReserve(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
limiter := NewInMemory()
|
|
policy := Policy{
|
|
Requests: 1,
|
|
Window: time.Hour,
|
|
Burst: 1,
|
|
}
|
|
|
|
first := limiter.Reserve("bucket-1", policy)
|
|
second := limiter.Reserve("bucket-1", policy)
|
|
otherBucket := limiter.Reserve("bucket-2", policy)
|
|
|
|
assert.True(t, first.Allowed)
|
|
assert.False(t, second.Allowed)
|
|
assert.Positive(t, second.RetryAfter)
|
|
assert.True(t, otherBucket.Allowed)
|
|
}
|
|
|
|
func TestInMemoryReserveResetsOnPolicyChange(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
limiter := NewInMemory()
|
|
|
|
initialPolicy := Policy{
|
|
Requests: 1,
|
|
Window: time.Hour,
|
|
Burst: 1,
|
|
}
|
|
updatedPolicy := Policy{
|
|
Requests: 2,
|
|
Window: time.Hour,
|
|
Burst: 2,
|
|
}
|
|
|
|
assert.True(t, limiter.Reserve("bucket-1", initialPolicy).Allowed)
|
|
assert.False(t, limiter.Reserve("bucket-1", initialPolicy).Allowed)
|
|
assert.True(t, limiter.Reserve("bucket-1", updatedPolicy).Allowed)
|
|
}
|