github.com/diamondburned/arikawa@v1.3.14/api/rate/rate_test.go (about)

     1  package rate
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"strconv"
     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  	headers := http.Header{}
    28  	headers.Set("X-RateLimit-Remaining", "0")
    29  	headers.Set("X-RateLimit-Reset",
    30  		strconv.FormatInt(time.Now().Add(time.Second*2).Unix(), 10))
    31  	headers.Set("Date", time.Now().Format(time.RFC850))
    32  
    33  	sent := time.Now()
    34  	mockRequest(t, l, "/guilds/99/channels", headers)
    35  	mockRequest(t, l, "/guilds/55/channels", headers)
    36  	mockRequest(t, l, "/guilds/66/channels", headers)
    37  
    38  	// call it again
    39  	mockRequest(t, l, "/guilds/99/channels", headers)
    40  	mockRequest(t, l, "/guilds/55/channels", headers)
    41  	mockRequest(t, l, "/guilds/66/channels", headers)
    42  
    43  	// We hit the same endpoint 2 times, so we should only be ratelimited 2
    44  	// second and always less than 4 seconds (unless you're on a stoneage
    45  	// computer or using swap or something...)
    46  	if time.Since(sent) >= time.Second && time.Since(sent) < time.Second*4 {
    47  		t.Log("OK", time.Since(sent))
    48  	} else {
    49  		t.Error("did not ratelimit correctly, got:", time.Since(sent))
    50  	}
    51  }
    52  
    53  // This test takes ~1 seconds to run
    54  func TestRatelimitGlobal(t *testing.T) {
    55  	l := NewLimiter("")
    56  
    57  	headers := http.Header{}
    58  	headers.Set("X-RateLimit-Global", "1.002")
    59  	// Reset for approx 1 seconds from now
    60  	headers.Set("Retry-After", "1000")
    61  
    62  	sent := time.Now()
    63  
    64  	// This should trigger a global ratelimit
    65  	mockRequest(t, l, "/guilds/99/channels", headers)
    66  	time.Sleep(time.Millisecond * 100)
    67  
    68  	// This shouldn't go through in less than 1 second
    69  	mockRequest(t, l, "/guilds/55/channels", headers)
    70  
    71  	if time.Since(sent) >= time.Second && time.Since(sent) < time.Second*2 {
    72  		t.Log("OK", time.Since(sent))
    73  	} else {
    74  		t.Error("did not ratelimit correctly, got:", time.Since(sent))
    75  	}
    76  }