46 lines
830 B
Go
46 lines
830 B
Go
package rate
|
|
|
|
import (
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
os.Exit(m.Run())
|
|
}
|
|
|
|
func TestRateLimit(t *testing.T) {
|
|
|
|
// 600 requests per 10 minutes
|
|
rate := NewLimit(600, time.Minute*10)
|
|
|
|
for i := 0; i < 600; i++ {
|
|
if err := rate.Try(); err != nil {
|
|
t.Log("should not have exceeded rate limit")
|
|
t.FailNow()
|
|
}
|
|
}
|
|
|
|
// Should fail
|
|
if err := rate.Try(); err == nil {
|
|
t.Log("should have exceeded rate limit")
|
|
t.FailNow()
|
|
}
|
|
|
|
// Pretend we had enough time to clear 5
|
|
rate.times = rate.times[5:]
|
|
for i := 0; i < 5; i++ {
|
|
if err := rate.Try(); err != nil {
|
|
t.Log("should not have exceeded rate limit")
|
|
t.FailNow()
|
|
}
|
|
}
|
|
|
|
// Make sure the error is checked correctly
|
|
if err := rate.Try(); !IsRateLimitExceededError(err) {
|
|
t.Log("error should be LimitExceededError")
|
|
t.FailNow()
|
|
}
|
|
}
|