github.com/koko1123/flow-go-1@v0.29.6/module/executiondatasync/pruner/pruner_test.go (about)

     1  package pruner_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/rs/zerolog"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/koko1123/flow-go-1/module/executiondatasync/pruner"
    12  	mocktracker "github.com/koko1123/flow-go-1/module/executiondatasync/tracker/mock"
    13  	"github.com/koko1123/flow-go-1/module/irrecoverable"
    14  	"github.com/koko1123/flow-go-1/module/metrics"
    15  	"github.com/koko1123/flow-go-1/utils/unittest"
    16  )
    17  
    18  func TestBasicPrune(t *testing.T) {
    19  	trackerStorage := new(mocktracker.Storage)
    20  	trackerStorage.On("GetFulfilledHeight").Return(uint64(0), nil).Once()
    21  	trackerStorage.On("GetPrunedHeight").Return(uint64(0), nil).Once()
    22  
    23  	pruner, err := pruner.NewPruner(
    24  		zerolog.Nop(),
    25  		metrics.NewNoopCollector(),
    26  		trackerStorage,
    27  		pruner.WithHeightRangeTarget(10),
    28  		pruner.WithThreshold(5),
    29  	)
    30  	require.NoError(t, err)
    31  	trackerStorage.AssertExpectations(t)
    32  
    33  	ctx, cancel := context.WithCancel(context.Background())
    34  	signalerCtx, errChan := irrecoverable.WithSignaler(ctx)
    35  
    36  	pruner.Start(signalerCtx)
    37  
    38  	pruned := make(chan struct{})
    39  	trackerStorage.On("PruneUpToHeight", uint64(6)).Return(func(height uint64) error {
    40  		close(pruned)
    41  		return nil
    42  	}).Once()
    43  
    44  	pruner.NotifyFulfilledHeight(16)
    45  	unittest.AssertClosesBefore(t, pruned, time.Second)
    46  	trackerStorage.AssertExpectations(t)
    47  
    48  	cancel()
    49  	<-pruner.Done()
    50  
    51  	select {
    52  	case err := <-errChan:
    53  		require.NoError(t, err)
    54  	default:
    55  	}
    56  }
    57  
    58  func TestInitialPrune(t *testing.T) {
    59  	trackerStorage := new(mocktracker.Storage)
    60  	trackerStorage.On("GetFulfilledHeight").Return(uint64(20), nil).Once()
    61  	trackerStorage.On("GetPrunedHeight").Return(uint64(0), nil).Once()
    62  
    63  	pruner, err := pruner.NewPruner(
    64  		zerolog.Nop(),
    65  		metrics.NewNoopCollector(),
    66  		trackerStorage,
    67  		pruner.WithHeightRangeTarget(10),
    68  		pruner.WithThreshold(5),
    69  	)
    70  	require.NoError(t, err)
    71  	trackerStorage.AssertExpectations(t)
    72  
    73  	ctx, cancel := context.WithCancel(context.Background())
    74  	signalerCtx, errChan := irrecoverable.WithSignaler(ctx)
    75  
    76  	pruned := make(chan struct{})
    77  	trackerStorage.On("PruneUpToHeight", uint64(10)).Return(func(height uint64) error {
    78  		close(pruned)
    79  		return nil
    80  	}).Once()
    81  
    82  	pruner.Start(signalerCtx)
    83  	unittest.AssertClosesBefore(t, pruned, time.Second)
    84  	trackerStorage.AssertExpectations(t)
    85  
    86  	cancel()
    87  	<-pruner.Done()
    88  
    89  	select {
    90  	case err := <-errChan:
    91  		require.NoError(t, err)
    92  	default:
    93  	}
    94  }
    95  
    96  func TestUpdateThreshold(t *testing.T) {
    97  	trackerStorage := new(mocktracker.Storage)
    98  	trackerStorage.On("GetFulfilledHeight").Return(uint64(15), nil).Once()
    99  	trackerStorage.On("GetPrunedHeight").Return(uint64(0), nil).Once()
   100  
   101  	pruner, err := pruner.NewPruner(
   102  		zerolog.Nop(),
   103  		metrics.NewNoopCollector(),
   104  		trackerStorage,
   105  		pruner.WithHeightRangeTarget(10),
   106  		pruner.WithThreshold(10),
   107  	)
   108  	require.NoError(t, err)
   109  	trackerStorage.AssertExpectations(t)
   110  
   111  	ctx, cancel := context.WithCancel(context.Background())
   112  	signalerCtx, errChan := irrecoverable.WithSignaler(ctx)
   113  
   114  	pruner.Start(signalerCtx)
   115  
   116  	pruned := make(chan struct{})
   117  	trackerStorage.On("PruneUpToHeight", uint64(5)).Return(func(height uint64) error {
   118  		close(pruned)
   119  		return nil
   120  	}).Once()
   121  
   122  	require.NoError(t, pruner.SetThreshold(4))
   123  	unittest.AssertClosesBefore(t, pruned, time.Second)
   124  	trackerStorage.AssertExpectations(t)
   125  
   126  	cancel()
   127  	<-pruner.Done()
   128  
   129  	select {
   130  	case err := <-errChan:
   131  		require.NoError(t, err)
   132  	default:
   133  	}
   134  }
   135  
   136  func TestUpdateHeightRangeTarget(t *testing.T) {
   137  	trackerStorage := new(mocktracker.Storage)
   138  	trackerStorage.On("GetFulfilledHeight").Return(uint64(10), nil).Once()
   139  	trackerStorage.On("GetPrunedHeight").Return(uint64(0), nil).Once()
   140  
   141  	pruner, err := pruner.NewPruner(
   142  		zerolog.Nop(),
   143  		metrics.NewNoopCollector(),
   144  		trackerStorage,
   145  		pruner.WithHeightRangeTarget(15),
   146  		pruner.WithThreshold(0),
   147  	)
   148  	require.NoError(t, err)
   149  	trackerStorage.AssertExpectations(t)
   150  
   151  	ctx, cancel := context.WithCancel(context.Background())
   152  	signalerCtx, errChan := irrecoverable.WithSignaler(ctx)
   153  
   154  	pruner.Start(signalerCtx)
   155  
   156  	pruned := make(chan struct{})
   157  	trackerStorage.On("PruneUpToHeight", uint64(5)).Return(func(height uint64) error {
   158  		close(pruned)
   159  		return nil
   160  	}).Once()
   161  
   162  	require.NoError(t, pruner.SetHeightRangeTarget(5))
   163  	unittest.AssertClosesBefore(t, pruned, time.Second)
   164  	trackerStorage.AssertExpectations(t)
   165  
   166  	cancel()
   167  	<-pruner.Done()
   168  
   169  	select {
   170  	case err := <-errChan:
   171  		require.NoError(t, err)
   172  	default:
   173  	}
   174  }