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

     1  package limit
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/alicebob/miniredis/v2"
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/shuguocloud/go-zero/core/stores/redis"
     9  	"github.com/shuguocloud/go-zero/core/stores/redis/redistest"
    10  )
    11  
    12  func TestPeriodLimit_Take(t *testing.T) {
    13  	testPeriodLimit(t)
    14  }
    15  
    16  func TestPeriodLimit_TakeWithAlign(t *testing.T) {
    17  	testPeriodLimit(t, Align())
    18  }
    19  
    20  func TestPeriodLimit_RedisUnavailable(t *testing.T) {
    21  	s, err := miniredis.Run()
    22  	assert.Nil(t, err)
    23  
    24  	const (
    25  		seconds = 1
    26  		total   = 100
    27  		quota   = 5
    28  	)
    29  	l := NewPeriodLimit(seconds, quota, redis.NewRedis(s.Addr(), redis.NodeType), "periodlimit")
    30  	s.Close()
    31  	val, err := l.Take("first")
    32  	assert.NotNil(t, err)
    33  	assert.Equal(t, 0, val)
    34  }
    35  
    36  func testPeriodLimit(t *testing.T, opts ...PeriodOption) {
    37  	store, clean, err := redistest.CreateRedis()
    38  	assert.Nil(t, err)
    39  	defer clean()
    40  
    41  	const (
    42  		seconds = 1
    43  		total   = 100
    44  		quota   = 5
    45  	)
    46  	l := NewPeriodLimit(seconds, quota, store, "periodlimit", opts...)
    47  	var allowed, hitQuota, overQuota int
    48  	for i := 0; i < total; i++ {
    49  		val, err := l.Take("first")
    50  		if err != nil {
    51  			t.Error(err)
    52  		}
    53  		switch val {
    54  		case Allowed:
    55  			allowed++
    56  		case HitQuota:
    57  			hitQuota++
    58  		case OverQuota:
    59  			overQuota++
    60  		default:
    61  			t.Error("unknown status")
    62  		}
    63  	}
    64  
    65  	assert.Equal(t, quota-1, allowed)
    66  	assert.Equal(t, 1, hitQuota)
    67  	assert.Equal(t, total-quota, overQuota)
    68  }