github.com/koko1123/flow-go-1@v0.29.6/utils/unittest/mocks/epoch_query.go (about)

     1  package mocks
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/assert"
     7  
     8  	"github.com/koko1123/flow-go-1/state/protocol"
     9  	"github.com/koko1123/flow-go-1/state/protocol/invalid"
    10  )
    11  
    12  // EpochQuery implements protocol.EpochQuery for testing purposes.
    13  type EpochQuery struct {
    14  	t         *testing.T
    15  	counter   uint64                    // represents the current epoch
    16  	byCounter map[uint64]protocol.Epoch // all epochs
    17  }
    18  
    19  func NewEpochQuery(t *testing.T, counter uint64, epochs ...protocol.Epoch) *EpochQuery {
    20  	mock := &EpochQuery{
    21  		t:         t,
    22  		counter:   counter,
    23  		byCounter: make(map[uint64]protocol.Epoch),
    24  	}
    25  
    26  	for _, epoch := range epochs {
    27  		mock.Add(epoch)
    28  	}
    29  
    30  	return mock
    31  }
    32  
    33  func (mock *EpochQuery) Current() protocol.Epoch {
    34  	return mock.byCounter[mock.counter]
    35  }
    36  
    37  func (mock *EpochQuery) Next() protocol.Epoch {
    38  	epoch, exists := mock.byCounter[mock.counter+1]
    39  	if !exists {
    40  		return invalid.NewEpoch(protocol.ErrNextEpochNotSetup)
    41  	}
    42  	return epoch
    43  }
    44  
    45  func (mock *EpochQuery) Previous() protocol.Epoch {
    46  	epoch, exists := mock.byCounter[mock.counter-1]
    47  	if !exists {
    48  		return invalid.NewEpoch(protocol.ErrNoPreviousEpoch)
    49  	}
    50  	return epoch
    51  }
    52  
    53  func (mock *EpochQuery) ByCounter(counter uint64) protocol.Epoch {
    54  	return mock.byCounter[counter]
    55  }
    56  
    57  func (mock *EpochQuery) Transition() {
    58  	mock.counter++
    59  }
    60  
    61  func (mock *EpochQuery) Add(epoch protocol.Epoch) {
    62  	counter, err := epoch.Counter()
    63  	assert.Nil(mock.t, err, "cannot add epoch with invalid counter")
    64  	mock.byCounter[counter] = epoch
    65  }