github.com/diamondburned/arikawa/v2@v2.1.0/api/rate/rate_test.go (about)

     1  package rate
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  	"testing"
     8  	"time"
     9  )
    10  
    11  // https://github.com/bwmarrin/discordgo/blob/master/ratelimit_test.go
    12  
    13  func mockRequest(t *testing.T, l *Limiter, path string, headers http.Header) {
    14  	if err := l.Acquire(context.Background(), path); err != nil {
    15  		t.Fatal("Failed to acquire lock:", err)
    16  	}
    17  
    18  	if err := l.Release(path, headers); err != nil {
    19  		t.Fatal("Failed to release lock:", err)
    20  	}
    21  }
    22  
    23  // This test takes ~2 seconds to run
    24  func TestRatelimitReset(t *testing.T) {
    25  	l := NewLimiter("")
    26  
    27  	const msToSec = time.Second / time.Millisecond
    28  
    29  	until := time.Now().Add(2 * time.Second)
    30  	reset := float64(until.UnixNano()/int64(time.Millisecond)) / float64(msToSec)
    31  
    32  	headers := http.Header{}
    33  	headers.Set("X-RateLimit-Remaining", "0")
    34  	headers.Set("X-RateLimit-Reset", fmt.Sprintf("%.3f", reset))
    35  	headers.Set("Date", time.Now().Format(time.RFC850))
    36  
    37  	sent := time.Now()
    38  	mockRequest(t, l, "/guilds/99/channels", headers)
    39  	mockRequest(t, l, "/guilds/55/channels", headers)
    40  	mockRequest(t, l, "/guilds/66/channels", headers)
    41  
    42  	// call it again
    43  	mockRequest(t, l, "/guilds/99/channels", headers)
    44  	mockRequest(t, l, "/guilds/55/channels", headers)
    45  	mockRequest(t, l, "/guilds/66/channels", headers)
    46  
    47  	// We hit the same endpoint 2 times, so we should only be ratelimited 2
    48  	// second and always less than 4 seconds (unless you're on a stoneage
    49  	// computer or using swap or something...)
    50  	if since := time.Since(sent); since >= time.Second && since < time.Second*4 {
    51  		t.Log("OK", since)
    52  	} else {
    53  		t.Error("did not ratelimit correctly, got:", since)
    54  	}
    55  }
    56  
    57  // This test takes ~1 seconds to run
    58  func TestRatelimitGlobal(t *testing.T) {
    59  	l := NewLimiter("")
    60  
    61  	headers := http.Header{}
    62  	headers.Set("X-RateLimit-Global", "1.002")
    63  	// Reset for approx 1 seconds from now
    64  	headers.Set("Retry-After", "1")
    65  
    66  	sent := time.Now()
    67  
    68  	// This should trigger a global ratelimit
    69  	mockRequest(t, l, "/guilds/99/channels", headers)
    70  	time.Sleep(time.Millisecond * 100)
    71  
    72  	// This shouldn't go through in less than 1 second
    73  	mockRequest(t, l, "/guilds/55/channels", headers)
    74  
    75  	if time.Since(sent) >= time.Second && time.Since(sent) < time.Second*2 {
    76  		t.Log("OK", time.Since(sent))
    77  	} else {
    78  		t.Error("did not ratelimit correctly, got:", time.Since(sent))
    79  	}
    80  }