github.com/peterdeka/grab@v2.0.0+incompatible/rate_limiter_test.go (about)

     1  package grab
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"os"
     8  	"testing"
     9  	"time"
    10  )
    11  
    12  // testRateLimiter is a naive rate limiter that limits throughput to r tokens
    13  // per second. The total number of tokens issued is tracked as n.
    14  type testRateLimiter struct {
    15  	r, n int
    16  }
    17  
    18  func NewLimiter(r int) RateLimiter {
    19  	return &testRateLimiter{r: r}
    20  }
    21  
    22  func (c *testRateLimiter) WaitN(ctx context.Context, n int) (err error) {
    23  	c.n += n
    24  	time.Sleep(
    25  		time.Duration(1.00 / float64(c.r) * float64(n) * float64(time.Second)))
    26  	return
    27  }
    28  
    29  func TestRateLimiter(t *testing.T) {
    30  	// download a 128 byte file, 8 bytes at a time, with a naive 512bps limiter
    31  	// should take > 250ms
    32  	filesize := 128
    33  	filename := ".testRateLimiter"
    34  	defer os.Remove(filename)
    35  
    36  	req, err := NewRequest(filename, fmt.Sprintf("%s?size=%d", ts.URL, filesize))
    37  	if err != nil {
    38  		t.Fatal(err)
    39  	}
    40  
    41  	// ensure multiple trips to the rate limiter by downloading 8 bytes at a time
    42  	req.BufferSize = 8
    43  
    44  	// limit to 512bps
    45  	lim := &testRateLimiter{r: 512}
    46  	req.RateLimiter = lim
    47  
    48  	resp := DefaultClient.Do(req)
    49  	if err = resp.Err(); err != nil {
    50  		t.Fatal(err)
    51  	}
    52  	testComplete(t, resp)
    53  	if lim.n != filesize {
    54  		t.Errorf("expected %d bytes to pass through limiter, got %d", filesize, lim.n)
    55  	}
    56  	if resp.Duration().Seconds() < 0.25 {
    57  		// BUG: this test can pass if the transfer was slow for unrelated reasons
    58  		t.Errorf("expected transfer to take >250ms, took %v", resp.Duration())
    59  	}
    60  }
    61  
    62  func ExampleRateLimiter() {
    63  	req, _ := NewRequest("", "http://www.golang-book.com/public/pdf/gobook.pdf")
    64  
    65  	// Attach a 1Mbps rate limiter, like the token bucket implementation from
    66  	// golang.org/x/time/rate.
    67  	req.RateLimiter = NewLimiter(1048576)
    68  
    69  	resp := DefaultClient.Do(req)
    70  	if err := resp.Err(); err != nil {
    71  		log.Fatal(err)
    72  	}
    73  }