github.com/blend/go-sdk@v1.20220411.3/ratelimiter/leaky_bucket_test.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package ratelimiter
     9  
    10  import (
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/blend/go-sdk/assert"
    15  )
    16  
    17  func TestLeakyBucket_Check(t *testing.T) {
    18  	it := assert.New(t)
    19  
    20  	rl := NewLeakyBucket(5, time.Second) // 5 actions per second
    21  
    22  	now := time.Now()
    23  
    24  	rl.Now = Clock(now, 0)
    25  	it.False(rl.Check("a"), "first call to `a` should pass")
    26  
    27  	rl.Now = Clock(now, 100*time.Millisecond)
    28  	it.False(rl.Check("b"), "first call to `b` should pass")
    29  
    30  	rl.Now = Clock(now, 200*time.Millisecond)
    31  	it.False(rl.Check("b"), "second call to `b` should pass")
    32  
    33  	rl.Now = Clock(now, 300*time.Millisecond)
    34  	it.False(rl.Check("b"), "third call to `b` should pass")
    35  
    36  	rl.Now = Clock(now, 400*time.Millisecond)
    37  	it.False(rl.Check("b"), "fourth call to `b` should pass")
    38  
    39  	rl.Now = Clock(now, 500*time.Millisecond)
    40  	it.False(rl.Check("a"), "second call to `a` in 500ms should pass")
    41  
    42  	rl.Now = Clock(now, 600*time.Millisecond)
    43  	it.False(rl.Check("a"), "third call to `a` in 600ms should pass")
    44  
    45  	rl.Now = Clock(now, 700*time.Millisecond)
    46  	it.False(rl.Check("a"), "fourth call to `a` in 700ms should pass")
    47  
    48  	rl.Now = Clock(now, 800*time.Millisecond)
    49  	it.True(rl.Check("a"), "fifth call to `a` in 800ms should fail")
    50  
    51  	rl.Now = Clock(now, 2000*time.Millisecond)
    52  	it.False(rl.Check("a"), "first call to `a` after pause should pass")
    53  
    54  	rl.Now = Clock(now, 2100*time.Millisecond)
    55  	it.False(rl.Check("b"), "first call to `b` after pause should pass")
    56  
    57  	rl.Now = Clock(now, 2200*time.Millisecond)
    58  	it.False(rl.Check("b"), "second call to `b` after pause should pass")
    59  
    60  	rl.Now = Clock(now, 2300*time.Millisecond)
    61  	it.False(rl.Check("b"), "third call to `b` after pause should pass")
    62  
    63  	rl.Now = Clock(now, 2400*time.Millisecond)
    64  	it.False(rl.Check("b"), "fourth call to `b` after pause should pass")
    65  
    66  	rl.Now = Clock(now, 2500*time.Millisecond)
    67  	it.False(rl.Check("a"), "second call to `a` after pause should pass")
    68  
    69  	rl.Now = Clock(now, 2600*time.Millisecond)
    70  	it.False(rl.Check("a"), "third call to `a` after pause should pass")
    71  
    72  	rl.Now = Clock(now, 2700*time.Millisecond)
    73  	it.False(rl.Check("a"), "fourth call to `a` after pause should pass")
    74  
    75  	rl.Now = Clock(now, 2800*time.Millisecond)
    76  	it.True(rl.Check("a"), "fifth call to `a` after pause should fail")
    77  }