wings/system/rate_test.go
Matthew Penner 3a738e44d6
run gofumpt
2022-02-23 15:02:19 -07:00

68 lines
1.3 KiB
Go

package system
import (
"testing"
"time"
. "github.com/franela/goblin"
)
func TestRate(t *testing.T) {
g := Goblin(t)
g.Describe("Rate", func() {
g.It("properly rate limits a bucket", func() {
r := NewRate(10, time.Millisecond*100)
for i := 0; i < 100; i++ {
ok := r.Try()
if i < 10 && !ok {
g.Failf("should not have allowed take on try %d", i)
} else if i >= 10 && ok {
g.Failf("should have blocked take on try %d", i)
}
}
})
g.It("handles rate limiting in chunks", func() {
var out []int
r := NewRate(12, time.Millisecond*10)
for i := 0; i < 100; i++ {
if i%20 == 0 {
// Give it time to recover.
time.Sleep(time.Millisecond * 10)
}
if r.Try() {
out = append(out, i)
}
}
g.Assert(len(out)).Equal(60)
g.Assert(out[0]).Equal(0)
g.Assert(out[12]).Equal(20)
g.Assert(out[len(out)-1]).Equal(91)
})
g.It("resets back to zero when called", func() {
r := NewRate(10, time.Second)
for i := 0; i < 100; i++ {
if i%10 == 0 {
r.Reset()
}
g.Assert(r.Try()).IsTrue()
}
g.Assert(r.Try()).IsFalse("final attempt should not allow taking")
})
})
}
func BenchmarkRate_Try(b *testing.B) {
r := NewRate(10, time.Millisecond*100)
b.ReportAllocs()
for i := 0; i < b.N; i++ {
r.Try()
}
}