vendor/github.com/go-kit/kit/ratelimit/token_bucket_test.go

   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 }