github.com/rudderlabs/rudder-go-kit@v0.30.0/throttling/memory_gcra_test.go (about)

     1  package throttling
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  func TestMemoryGCRA(t *testing.T) {
    12  	t.Run("burst and cost greater than rate", func(t *testing.T) {
    13  		l := &gcra{}
    14  
    15  		burst := int64(5)
    16  		rate := int64(1)
    17  		period := int64(1)
    18  
    19  		allowed, _, err := l.limit(context.Background(), "key", burst+rate, burst, rate, period)
    20  		require.NoError(t, err)
    21  		require.True(t, allowed, "it should be able to fill the bucket (burst)")
    22  
    23  		// next request should be allowed after 5 seconds
    24  		start := time.Now()
    25  
    26  		require.Eventually(t, func() bool {
    27  			allowed, _, err := l.limit(context.Background(), "key", burst, burst, rate, period)
    28  			if err != nil {
    29  				t.Logf("Memory GCRA error: %v", err)
    30  				return false
    31  			}
    32  			return allowed
    33  		}, 10*time.Second, 1*time.Second, "next request should be eventually allowed")
    34  
    35  		require.GreaterOrEqual(t, time.Since(start).Seconds(), 5.0, "next request should be allowed after 5 seconds")
    36  	})
    37  }