github.com/raamsri/grab/v3@v3.0.2/rate_limiter_test.go (about)

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