1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2024-12-12 08:25:10 +00:00
arikawa/api/rate/rate_test.go

81 lines
2.2 KiB
Go
Raw Normal View History

2020-01-08 07:10:37 +00:00
package rate
import (
"context"
2020-10-28 22:23:47 +00:00
"fmt"
2020-01-08 07:10:37 +00:00
"net/http"
"testing"
"time"
)
// https://github.com/bwmarrin/discordgo/blob/master/ratelimit_test.go
func mockRequest(t *testing.T, l *Limiter, path string, headers http.Header) {
if err := l.Acquire(context.Background(), path); err != nil {
t.Fatal("Failed to acquire lock:", err)
}
if err := l.Release(path, headers); err != nil {
t.Fatal("Failed to release lock:", err)
}
}
// This test takes ~2 seconds to run
func TestRatelimitReset(t *testing.T) {
l := NewLimiter("")
2020-01-08 07:10:37 +00:00
2020-10-28 22:23:47 +00:00
const msToSec = time.Second / time.Millisecond
until := time.Now().Add(2 * time.Second)
reset := float64(until.UnixNano()/int64(time.Millisecond)) / float64(msToSec)
2020-01-08 07:10:37 +00:00
headers := http.Header{}
headers.Set("X-RateLimit-Remaining", "0")
2020-10-28 22:23:47 +00:00
headers.Set("X-RateLimit-Reset", fmt.Sprintf("%.3f", reset))
2020-01-08 07:10:37 +00:00
headers.Set("Date", time.Now().Format(time.RFC850))
sent := time.Now()
mockRequest(t, l, "/guilds/99/channels", headers)
mockRequest(t, l, "/guilds/55/channels", headers)
mockRequest(t, l, "/guilds/66/channels", headers)
// call it again
mockRequest(t, l, "/guilds/99/channels", headers)
mockRequest(t, l, "/guilds/55/channels", headers)
mockRequest(t, l, "/guilds/66/channels", headers)
// We hit the same endpoint 2 times, so we should only be ratelimited 2
// second and always less than 4 seconds (unless you're on a stoneage
// computer or using swap or something...)
2020-10-28 22:23:47 +00:00
if since := time.Since(sent); since >= time.Second && since < time.Second*4 {
t.Log("OK", since)
2020-01-08 07:10:37 +00:00
} else {
2020-10-28 22:23:47 +00:00
t.Error("did not ratelimit correctly, got:", since)
2020-01-08 07:10:37 +00:00
}
}
// This test takes ~1 seconds to run
func TestRatelimitGlobal(t *testing.T) {
l := NewLimiter("")
2020-01-08 07:10:37 +00:00
headers := http.Header{}
2020-01-09 05:24:45 +00:00
headers.Set("X-RateLimit-Global", "1.002")
2020-01-08 07:10:37 +00:00
// Reset for approx 1 seconds from now
2020-10-28 22:23:47 +00:00
headers.Set("Retry-After", "1")
2020-01-08 07:10:37 +00:00
sent := time.Now()
// This should trigger a global ratelimit
mockRequest(t, l, "/guilds/99/channels", headers)
time.Sleep(time.Millisecond * 100)
// This shouldn't go through in less than 1 second
mockRequest(t, l, "/guilds/55/channels", headers)
if time.Since(sent) >= time.Second && time.Since(sent) < time.Second*2 {
t.Log("OK", time.Since(sent))
} else {
2020-05-16 21:14:49 +00:00
t.Error("did not ratelimit correctly, got:", time.Since(sent))
2020-01-08 07:10:37 +00:00
}
}