github.com/koko1123/flow-go-1@v0.29.6/storage/badger/operation/dkg_test.go (about)

     1  package operation
     2  
     3  import (
     4  	"math/rand"
     5  	"testing"
     6  
     7  	"github.com/dgraph-io/badger/v3"
     8  	"github.com/stretchr/testify/assert"
     9  
    10  	"github.com/koko1123/flow-go-1/model/encodable"
    11  	"github.com/koko1123/flow-go-1/model/flow"
    12  	"github.com/koko1123/flow-go-1/storage"
    13  	"github.com/koko1123/flow-go-1/utils/unittest"
    14  )
    15  
    16  // TestInsertMyDKGPrivateInfo_StoreRetrieve tests writing and reading private DKG info.
    17  func TestMyBeaconPrivateKey_StoreRetrieve(t *testing.T) {
    18  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    19  
    20  		t.Run("should return error not found when not stored", func(t *testing.T) {
    21  			var stored encodable.RandomBeaconPrivKey
    22  			err := db.View(RetrieveMyBeaconPrivateKey(1, &stored))
    23  			assert.ErrorIs(t, err, storage.ErrNotFound)
    24  		})
    25  
    26  		t.Run("should be able to store and read", func(t *testing.T) {
    27  			epochCounter := rand.Uint64()
    28  			info := unittest.RandomBeaconPriv()
    29  
    30  			// should be able to store
    31  			err := db.Update(InsertMyBeaconPrivateKey(epochCounter, info))
    32  			assert.NoError(t, err)
    33  
    34  			// should be able to read
    35  			var stored encodable.RandomBeaconPrivKey
    36  			err = db.View(RetrieveMyBeaconPrivateKey(epochCounter, &stored))
    37  			assert.NoError(t, err)
    38  			assert.Equal(t, info, &stored)
    39  
    40  			// should fail to read other epoch counter
    41  			err = db.View(RetrieveMyBeaconPrivateKey(rand.Uint64(), &stored))
    42  			assert.ErrorIs(t, err, storage.ErrNotFound)
    43  		})
    44  	})
    45  }
    46  
    47  // TestDKGStartedForEpoch tests setting the DKG-started flag.
    48  func TestDKGStartedForEpoch(t *testing.T) {
    49  
    50  	t.Run("reading when unset should return false", func(t *testing.T) {
    51  		unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    52  			var started bool
    53  			err := db.View(RetrieveDKGStartedForEpoch(1, &started))
    54  			assert.NoError(t, err)
    55  			assert.False(t, started)
    56  		})
    57  	})
    58  
    59  	t.Run("should be able to set flag to true", func(t *testing.T) {
    60  		unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    61  			epochCounter := rand.Uint64()
    62  
    63  			// set the flag, ensure no error
    64  			err := db.Update(InsertDKGStartedForEpoch(epochCounter))
    65  			assert.NoError(t, err)
    66  
    67  			// read the flag, should be true now
    68  			var started bool
    69  			err = db.View(RetrieveDKGStartedForEpoch(epochCounter, &started))
    70  			assert.NoError(t, err)
    71  			assert.True(t, started)
    72  
    73  			// read the flag for a different epoch, should be false
    74  			err = db.View(RetrieveDKGStartedForEpoch(epochCounter+1, &started))
    75  			assert.NoError(t, err)
    76  			assert.False(t, started)
    77  		})
    78  	})
    79  }
    80  
    81  func TestDKGEndStateForEpoch(t *testing.T) {
    82  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    83  		epochCounter := rand.Uint64()
    84  
    85  		// should be able to write end state
    86  		endState := flow.DKGEndStateSuccess
    87  		err := db.Update(InsertDKGEndStateForEpoch(epochCounter, endState))
    88  		assert.NoError(t, err)
    89  
    90  		// should be able to read end state
    91  		var readEndState flow.DKGEndState
    92  		err = db.View(RetrieveDKGEndStateForEpoch(epochCounter, &readEndState))
    93  		assert.NoError(t, err)
    94  		assert.Equal(t, endState, readEndState)
    95  
    96  		// attempting to overwrite should error
    97  		err = db.Update(InsertDKGEndStateForEpoch(epochCounter, flow.DKGEndStateDKGFailure))
    98  		assert.ErrorIs(t, err, storage.ErrAlreadyExists)
    99  	})
   100  }