github.com/ethersphere/bee/v2@v2.2.0/pkg/ratelimit/ratelimit_test.go (about)

     1  // Copyright 2021 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package ratelimit_test
     6  
     7  import (
     8  	"context"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/ethersphere/bee/v2/pkg/ratelimit"
    13  )
    14  
    15  func TestRateLimit(t *testing.T) {
    16  	t.Parallel()
    17  
    18  	var (
    19  		key1  = "test1"
    20  		key2  = "test2"
    21  		rate  = time.Second
    22  		burst = 10
    23  	)
    24  
    25  	limiter := ratelimit.New(rate, burst)
    26  
    27  	if !limiter.Allow(key1, burst) {
    28  		t.Fatal("want allowed")
    29  	}
    30  
    31  	if limiter.Allow(key1, burst) {
    32  		t.Fatalf("want not allowed")
    33  	}
    34  
    35  	limiter.Clear(key1)
    36  
    37  	if !limiter.Allow(key1, burst) {
    38  		t.Fatal("want allowed")
    39  	}
    40  
    41  	if !limiter.Allow(key2, burst) {
    42  		t.Fatal("want allowed")
    43  	}
    44  }
    45  
    46  func TestWait(t *testing.T) {
    47  	t.Parallel()
    48  
    49  	var (
    50  		key   = "test"
    51  		rate  = time.Second
    52  		burst = 4
    53  	)
    54  
    55  	limiter := ratelimit.New(rate, burst)
    56  
    57  	if !limiter.Allow(key, 1) {
    58  		t.Fatal("want allowed")
    59  	}
    60  
    61  	waitDur, err := limiter.Wait(context.Background(), key, 1)
    62  	if err != nil {
    63  		t.Fatalf("got err %v", err)
    64  	}
    65  
    66  	if waitDur != 0 {
    67  		t.Fatalf("expected the limiter to NOT wait, got %s", waitDur)
    68  	}
    69  
    70  	waitDur, err = limiter.Wait(context.Background(), key, burst)
    71  	if err != nil {
    72  		t.Fatalf("got err %v", err)
    73  	}
    74  
    75  	if waitDur < rate {
    76  		t.Fatalf("expected the limiter to wait at least %s, got %s", rate, waitDur)
    77  	}
    78  }