github.com/shuguocloud/go-zero@v1.3.0/core/limit/tokenlimit_test.go (about)

     1  package limit
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/alicebob/miniredis/v2"
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/shuguocloud/go-zero/core/logx"
    10  	"github.com/shuguocloud/go-zero/core/stores/redis"
    11  	"github.com/shuguocloud/go-zero/core/stores/redis/redistest"
    12  )
    13  
    14  func init() {
    15  	logx.Disable()
    16  }
    17  
    18  func TestTokenLimit_Rescue(t *testing.T) {
    19  	s, err := miniredis.Run()
    20  	assert.Nil(t, err)
    21  
    22  	const (
    23  		total = 100
    24  		rate  = 5
    25  		burst = 10
    26  	)
    27  	l := NewTokenLimiter(rate, burst, redis.New(s.Addr()), "tokenlimit")
    28  	s.Close()
    29  
    30  	var allowed int
    31  	for i := 0; i < total; i++ {
    32  		time.Sleep(time.Second / time.Duration(total))
    33  		if i == total>>1 {
    34  			assert.Nil(t, s.Restart())
    35  		}
    36  		if l.Allow() {
    37  			allowed++
    38  		}
    39  
    40  		// make sure start monitor more than once doesn't matter
    41  		l.startMonitor()
    42  	}
    43  
    44  	assert.True(t, allowed >= burst+rate)
    45  }
    46  
    47  func TestTokenLimit_Take(t *testing.T) {
    48  	store, clean, err := redistest.CreateRedis()
    49  	assert.Nil(t, err)
    50  	defer clean()
    51  
    52  	const (
    53  		total = 100
    54  		rate  = 5
    55  		burst = 10
    56  	)
    57  	l := NewTokenLimiter(rate, burst, store, "tokenlimit")
    58  	var allowed int
    59  	for i := 0; i < total; i++ {
    60  		time.Sleep(time.Second / time.Duration(total))
    61  		if l.Allow() {
    62  			allowed++
    63  		}
    64  	}
    65  
    66  	assert.True(t, allowed >= burst+rate)
    67  }
    68  
    69  func TestTokenLimit_TakeBurst(t *testing.T) {
    70  	store, clean, err := redistest.CreateRedis()
    71  	assert.Nil(t, err)
    72  	defer clean()
    73  
    74  	const (
    75  		total = 100
    76  		rate  = 5
    77  		burst = 10
    78  	)
    79  	l := NewTokenLimiter(rate, burst, store, "tokenlimit")
    80  	var allowed int
    81  	for i := 0; i < total; i++ {
    82  		if l.Allow() {
    83  			allowed++
    84  		}
    85  	}
    86  
    87  	assert.True(t, allowed >= burst)
    88  }