github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/kv/store_limiter_test.go (about)

     1  package kv_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/golang/mock/gomock"
     9  	"github.com/treeverse/lakefs/pkg/kv"
    10  	"github.com/treeverse/lakefs/pkg/kv/mock"
    11  )
    12  
    13  type traceLimiter struct {
    14  	Count int
    15  }
    16  
    17  func (t *traceLimiter) Take() time.Time {
    18  	t.Count++
    19  	return time.Now()
    20  }
    21  
    22  // TestStoreLimiter verify limiter is used for all kv store operations
    23  func TestStoreLimiter(t *testing.T) {
    24  	ctx := context.Background()
    25  	ctrl := gomock.NewController(t)
    26  	store := mock.NewMockStore(ctrl)
    27  
    28  	rec := store.EXPECT()
    29  	rec.Scan(gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil, nil)
    30  	rec.Get(gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil, nil)
    31  	rec.Delete(gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
    32  	rec.Set(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
    33  	rec.SetIf(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
    34  	rec.Close().Times(1)
    35  
    36  	limiter := &traceLimiter{}
    37  	sl := kv.NewStoreLimiter(store, limiter)
    38  	defer sl.Close()
    39  	_, _ = sl.Scan(ctx, nil, kv.ScanOptions{})
    40  	_, _ = sl.Get(ctx, nil, nil)
    41  	_ = sl.Delete(ctx, nil, nil)
    42  	_ = sl.Set(ctx, nil, nil, nil)
    43  	_ = sl.SetIf(ctx, nil, nil, nil, nil)
    44  
    45  	const expectedCalls = 5
    46  	if expectedCalls != limiter.Count {
    47  		t.Fatalf("Limiter called=%d, expected=%d", limiter.Count, expectedCalls)
    48  	}
    49  }