github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/framework/ratelimit/memrate/memorystore_test.go (about)

     1  package memrate
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	lru "github.com/hashicorp/golang-lru"
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/unionj-cloud/go-doudou/framework/ratelimit"
     9  	"reflect"
    10  	"testing"
    11  )
    12  
    13  func TestMemoryStore_addKey(t *testing.T) {
    14  	type fields struct {
    15  		keys      map[string]ratelimit.Limiter
    16  		limiterFn LimiterFn
    17  	}
    18  	type args struct {
    19  		key string
    20  	}
    21  	limiter := NewLimiter(1, 3)
    22  	tests := []struct {
    23  		name   string
    24  		fields fields
    25  		args   args
    26  		want   ratelimit.Limiter
    27  	}{
    28  		{
    29  			name: "",
    30  			fields: fields{
    31  				keys: make(map[string]ratelimit.Limiter),
    32  				limiterFn: func(ctx context.Context, store *MemoryStore, key string) ratelimit.Limiter {
    33  					return NewLimiter(1, 3)
    34  				},
    35  			},
    36  			args: args{
    37  				key: "192.168.1.6:8080",
    38  			},
    39  			want: limiter,
    40  		},
    41  	}
    42  	for _, tt := range tests {
    43  		t.Run(tt.name, func(t *testing.T) {
    44  			keys, _ := lru.New(256)
    45  			store := &MemoryStore{
    46  				keys:      keys,
    47  				limiterFn: tt.fields.limiterFn,
    48  			}
    49  			if got := store.addKeyCtx(context.Background(), tt.args.key); !reflect.DeepEqual(got, tt.want) {
    50  				t.Errorf("addKey() = %v, want %v", got, tt.want)
    51  			}
    52  		})
    53  	}
    54  }
    55  
    56  func TestMemoryStore_GetLimiter(t *testing.T) {
    57  	type fields struct {
    58  		keys      map[string]ratelimit.Limiter
    59  		limiterFn LimiterFn
    60  	}
    61  	type args struct {
    62  		key string
    63  	}
    64  	limiter := NewLimiter(1, 3)
    65  	key := "192.168.1.6:8080"
    66  	keys, _ := lru.New(256)
    67  	store := &MemoryStore{
    68  		keys: keys,
    69  		limiterFn: func(ctx context.Context, store *MemoryStore, key string) ratelimit.Limiter {
    70  			return NewLimiter(1, 3)
    71  		},
    72  	}
    73  	assert.Equal(t, limiter, store.GetLimiter(key))
    74  	assert.Equal(t, limiter, store.GetLimiter(key))
    75  }
    76  
    77  func TestMemoryStore_DeleteKey(t *testing.T) {
    78  	key := "192.168.1.6:8080"
    79  	keys, _ := lru.New(256)
    80  	store := &MemoryStore{
    81  		keys: keys,
    82  		limiterFn: func(ctx context.Context, store *MemoryStore, key string) ratelimit.Limiter {
    83  			return NewLimiter(1, 3)
    84  		},
    85  	}
    86  	store.addKeyCtx(context.Background(), key)
    87  	store.addKeyCtx(context.Background(), key)
    88  	if exists := store.keys.Contains(key); !exists {
    89  		t.Error("key should exists")
    90  	}
    91  	store.DeleteKey(key)
    92  	if exists := store.keys.Contains(key); exists {
    93  		t.Error("key should not exists")
    94  	}
    95  }
    96  
    97  func TestNewMemoryStore(t *testing.T) {
    98  	type args struct {
    99  		fn   LimiterFn
   100  		opts []MemoryStoreOption
   101  	}
   102  	tests := []struct {
   103  		name string
   104  		args args
   105  	}{
   106  		{
   107  			name: "",
   108  			args: args{
   109  				fn: func(ctx context.Context, store *MemoryStore, key string) ratelimit.Limiter {
   110  					return NewLimiter(1, 3)
   111  				},
   112  				opts: nil,
   113  			},
   114  		},
   115  		{
   116  			name: "",
   117  			args: args{
   118  				fn: func(ctx context.Context, store *MemoryStore, key string) ratelimit.Limiter {
   119  					return NewLimiter(1, 3)
   120  				},
   121  				opts: []MemoryStoreOption{
   122  					WithMaxKeys(100),
   123  					WithOnEvicted(func(key interface{}, value interface{}) {
   124  						fmt.Println(key, value)
   125  					}),
   126  				},
   127  			},
   128  		},
   129  	}
   130  	for _, tt := range tests {
   131  		t.Run(tt.name, func(t *testing.T) {
   132  			if got := NewMemoryStore(tt.args.fn, tt.args.opts...); got == nil {
   133  				t.Error("NewMemoryStore() = nil")
   134  			}
   135  		})
   136  	}
   137  }