gopkg.in/cavaliercoder/grab.v2@v2.0.0/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 }