github.com/ulule/limiter/v3@v3.11.3-0.20230613131926-4cb9c1da4633/drivers/store/tests/tests.go (about)

     1  package tests
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/ulule/limiter/v3"
    12  )
    13  
    14  // TestStoreSequentialAccess verify that store works as expected with a sequential access.
    15  func TestStoreSequentialAccess(t *testing.T, store limiter.Store) {
    16  	is := require.New(t)
    17  	ctx := context.Background()
    18  
    19  	limiter := limiter.New(store, limiter.Rate{
    20  		Limit:  3,
    21  		Period: time.Minute,
    22  	})
    23  
    24  	// Check counter increment.
    25  	{
    26  		for i := 1; i <= 6; i++ {
    27  
    28  			if i <= 3 {
    29  
    30  				lctx, err := limiter.Peek(ctx, "foo")
    31  				is.NoError(err)
    32  				is.NotZero(lctx)
    33  				is.Equal(int64(3-(i-1)), lctx.Remaining)
    34  				is.False(lctx.Reached)
    35  
    36  			}
    37  
    38  			lctx, err := limiter.Get(ctx, "foo")
    39  			is.NoError(err)
    40  			is.NotZero(lctx)
    41  
    42  			if i <= 3 {
    43  
    44  				is.Equal(int64(3), lctx.Limit)
    45  				is.Equal(int64(3-i), lctx.Remaining)
    46  				is.True((lctx.Reset - time.Now().Unix()) <= 60)
    47  				is.False(lctx.Reached)
    48  
    49  				lctx, err = limiter.Peek(ctx, "foo")
    50  				is.NoError(err)
    51  				is.Equal(int64(3-i), lctx.Remaining)
    52  				is.False(lctx.Reached)
    53  
    54  			} else {
    55  
    56  				is.Equal(int64(3), lctx.Limit)
    57  				is.Equal(int64(0), lctx.Remaining)
    58  				is.True((lctx.Reset - time.Now().Unix()) <= 60)
    59  				is.True(lctx.Reached)
    60  
    61  			}
    62  		}
    63  	}
    64  
    65  	// Check counter reset.
    66  	{
    67  		lctx, err := limiter.Peek(ctx, "foo")
    68  		is.NoError(err)
    69  		is.NotZero(lctx)
    70  
    71  		is.Equal(int64(3), lctx.Limit)
    72  		is.Equal(int64(0), lctx.Remaining)
    73  		is.True((lctx.Reset - time.Now().Unix()) <= 60)
    74  		is.True(lctx.Reached)
    75  
    76  		lctx, err = limiter.Reset(ctx, "foo")
    77  		is.NoError(err)
    78  		is.NotZero(lctx)
    79  
    80  		is.Equal(int64(3), lctx.Limit)
    81  		is.Equal(int64(3), lctx.Remaining)
    82  		is.True((lctx.Reset - time.Now().Unix()) <= 60)
    83  		is.False(lctx.Reached)
    84  
    85  		lctx, err = limiter.Peek(ctx, "foo")
    86  		is.NoError(err)
    87  		is.NotZero(lctx)
    88  
    89  		is.Equal(int64(3), lctx.Limit)
    90  		is.Equal(int64(3), lctx.Remaining)
    91  		is.True((lctx.Reset - time.Now().Unix()) <= 60)
    92  		is.False(lctx.Reached)
    93  
    94  		lctx, err = limiter.Get(ctx, "foo")
    95  		is.NoError(err)
    96  		is.NotZero(lctx)
    97  
    98  		lctx, err = limiter.Reset(ctx, "foo")
    99  		is.NoError(err)
   100  		is.NotZero(lctx)
   101  
   102  		is.Equal(int64(3), lctx.Limit)
   103  		is.Equal(int64(3), lctx.Remaining)
   104  		is.True((lctx.Reset - time.Now().Unix()) <= 60)
   105  		is.False(lctx.Reached)
   106  
   107  		lctx, err = limiter.Reset(ctx, "foo")
   108  		is.NoError(err)
   109  		is.NotZero(lctx)
   110  
   111  		is.Equal(int64(3), lctx.Limit)
   112  		is.Equal(int64(3), lctx.Remaining)
   113  		is.True((lctx.Reset - time.Now().Unix()) <= 60)
   114  		is.False(lctx.Reached)
   115  	}
   116  }
   117  
   118  // TestStoreConcurrentAccess verify that store works as expected with a concurrent access.
   119  func TestStoreConcurrentAccess(t *testing.T, store limiter.Store) {
   120  	is := require.New(t)
   121  	ctx := context.Background()
   122  
   123  	limiter := limiter.New(store, limiter.Rate{
   124  		Limit:  100000,
   125  		Period: 10 * time.Second,
   126  	})
   127  
   128  	goroutines := 500
   129  	ops := 500
   130  
   131  	wg := &sync.WaitGroup{}
   132  	wg.Add(goroutines)
   133  	for i := 0; i < goroutines; i++ {
   134  		go func(i int) {
   135  			for j := 0; j < ops; j++ {
   136  				lctx, err := limiter.Get(ctx, "foo")
   137  				is.NoError(err)
   138  				is.NotZero(lctx)
   139  			}
   140  			wg.Done()
   141  		}(i)
   142  	}
   143  	wg.Wait()
   144  }
   145  
   146  // BenchmarkStoreSequentialAccess executes a benchmark against a store without parallel setting.
   147  func BenchmarkStoreSequentialAccess(b *testing.B, store limiter.Store) {
   148  	ctx := context.Background()
   149  
   150  	instance := limiter.New(store, limiter.Rate{
   151  		Limit:  100000,
   152  		Period: 10 * time.Second,
   153  	})
   154  
   155  	b.ResetTimer()
   156  	for i := 0; i < b.N; i++ {
   157  		_, _ = instance.Get(ctx, "foo")
   158  	}
   159  }
   160  
   161  // BenchmarkStoreConcurrentAccess executes a benchmark against a store with parallel setting.
   162  func BenchmarkStoreConcurrentAccess(b *testing.B, store limiter.Store) {
   163  	ctx := context.Background()
   164  
   165  	instance := limiter.New(store, limiter.Rate{
   166  		Limit:  100000,
   167  		Period: 10 * time.Second,
   168  	})
   169  
   170  	b.ResetTimer()
   171  	b.RunParallel(func(pb *testing.PB) {
   172  		for pb.Next() {
   173  			_, _ = instance.Get(ctx, "foo")
   174  		}
   175  	})
   176  }