github.com/cilium/cilium@v1.16.2/pkg/api/helpers/rate_limit_test.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Cilium
     3  
     4  package helpers
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/cilium/cilium/pkg/api/metrics/mock"
    14  )
    15  
    16  func TestRateLimitBurst(t *testing.T) {
    17  	metricsAPI := mock.NewMockMetrics()
    18  	limiter := NewAPILimiter(metricsAPI, 1, 10)
    19  	require.NotNil(t, limiter)
    20  
    21  	// Exhaust bucket (rate limit should not kick in)
    22  	for i := 0; i < 10; i++ {
    23  		limiter.Limit(context.TODO(), "test")
    24  	}
    25  	require.Equal(t, time.Duration(0), metricsAPI.RateLimit("test"))
    26  
    27  	// Rate limit should now kick in (use an expired context to avoid waiting 1sec)
    28  	ctx, cancel := context.WithTimeout(context.TODO(), time.Microsecond)
    29  	defer cancel()
    30  	limiter.Limit(ctx, "test")
    31  	require.NotEqual(t, time.Duration(0), metricsAPI.RateLimit("test"))
    32  }
    33  
    34  func TestRateLimitWait(t *testing.T) {
    35  	metricsAPI := mock.NewMockMetrics()
    36  	limiter := NewAPILimiter(metricsAPI, 100, 1)
    37  	require.NotNil(t, limiter)
    38  
    39  	// Exhaust bucket
    40  	limiter.Limit(context.TODO(), "test")
    41  	require.Equal(t, time.Duration(0), metricsAPI.RateLimit("test"))
    42  
    43  	// Hit rate limit 15 times. The bucket refill rate is 100 per second,
    44  	// meaning we expect this to take around 15 * 10 = 150 milliseconds
    45  	start := time.Now()
    46  	for i := 0; i < 15; i++ {
    47  		limiter.Limit(context.TODO(), "test")
    48  	}
    49  	measured := time.Since(start)
    50  
    51  	// Measured duration should be approximately the accounted duration
    52  	accounted := metricsAPI.RateLimit("test")
    53  	if measured > 2*accounted {
    54  		// We allow the wait to be up to 2x larger than the expected wait time
    55  		// to avoid flaky tests. If you are reading this because this test has
    56  		// been flaky despite the 100% margin of error, my recommendation
    57  		// is to disable this check by replacing the c.Errorf below with c.Logf
    58  		t.Errorf("waited longer than expected (expected %s (+/-100%%), measured %s)", accounted, measured)
    59  	}
    60  }