Source file src/github.com/go-kit/kit/ratelimit/token_bucket_test.go

Documentation: github.com/go-kit/kit/ratelimit

     1  package ratelimit_test
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"testing"
     7  	"time"
     8  
     9  	"golang.org/x/time/rate"
    10  
    11  	"github.com/go-kit/kit/endpoint"
    12  	"github.com/go-kit/kit/ratelimit"
    13  )
    14  
    15  var nopEndpoint = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
    16  
    17  func TestXRateErroring(t *testing.T) {
    18  	limit := rate.NewLimiter(rate.Every(time.Minute), 1)
    19  	testSuccessThenFailure(
    20  		t,
    21  		ratelimit.NewErroringLimiter(limit)(nopEndpoint),
    22  		ratelimit.ErrLimited.Error())
    23  }
    24  
    25  func TestXRateDelaying(t *testing.T) {
    26  	limit := rate.NewLimiter(rate.Every(time.Minute), 1)
    27  	testSuccessThenFailure(
    28  		t,
    29  		ratelimit.NewDelayingLimiter(limit)(nopEndpoint),
    30  		"exceed context deadline")
    31  }
    32  
    33  func testSuccessThenFailure(t *testing.T, e endpoint.Endpoint, failContains string) {
    34  	ctx, cxl := context.WithTimeout(context.Background(), 500*time.Millisecond)
    35  	defer cxl()
    36  
    37  	// First request should succeed.
    38  	if _, err := e(ctx, struct{}{}); err != nil {
    39  		t.Errorf("unexpected: %v\n", err)
    40  	}
    41  
    42  	// Next request should fail.
    43  	if _, err := e(ctx, struct{}{}); !strings.Contains(err.Error(), failContains) {
    44  		t.Errorf("expected `%s`: %v\n", failContains, err)
    45  	}
    46  }
    47  

View as plain text