github.com/lingyao2333/mo-zero@v1.4.1/core/limit/periodlimit_test.go (about)

     1  package limit
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/alicebob/miniredis/v2"
     7  	"github.com/lingyao2333/mo-zero/core/stores/redis"
     8  	"github.com/lingyao2333/mo-zero/core/stores/redis/redistest"
     9  	"github.com/stretchr/testify/assert"
    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  		quota   = 5
    27  	)
    28  	l := NewPeriodLimit(seconds, quota, redis.New(s.Addr()), "periodlimit")
    29  	s.Close()
    30  	val, err := l.Take("first")
    31  	assert.NotNil(t, err)
    32  	assert.Equal(t, 0, val)
    33  }
    34  
    35  func testPeriodLimit(t *testing.T, opts ...PeriodOption) {
    36  	store, clean, err := redistest.CreateRedis()
    37  	assert.Nil(t, err)
    38  	defer clean()
    39  
    40  	const (
    41  		seconds = 1
    42  		total   = 100
    43  		quota   = 5
    44  	)
    45  	l := NewPeriodLimit(seconds, quota, store, "periodlimit", opts...)
    46  	var allowed, hitQuota, overQuota int
    47  	for i := 0; i < total; i++ {
    48  		val, err := l.Take("first")
    49  		if err != nil {
    50  			t.Error(err)
    51  		}
    52  		switch val {
    53  		case Allowed:
    54  			allowed++
    55  		case HitQuota:
    56  			hitQuota++
    57  		case OverQuota:
    58  			overQuota++
    59  		default:
    60  			t.Error("unknown status")
    61  		}
    62  	}
    63  
    64  	assert.Equal(t, quota-1, allowed)
    65  	assert.Equal(t, 1, hitQuota)
    66  	assert.Equal(t, total-quota, overQuota)
    67  }
    68  
    69  func TestQuotaFull(t *testing.T) {
    70  	s, err := miniredis.Run()
    71  	assert.Nil(t, err)
    72  
    73  	l := NewPeriodLimit(1, 1, redis.New(s.Addr()), "periodlimit")
    74  	val, err := l.Take("first")
    75  	assert.Nil(t, err)
    76  	assert.Equal(t, HitQuota, val)
    77  }