github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/storage/badger/protocol_kv_store_test.go (about)

     1  package badger
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/dgraph-io/badger/v2"
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/require"
     9  
    10  	"github.com/onflow/flow-go/model/flow"
    11  	"github.com/onflow/flow-go/module/metrics"
    12  	"github.com/onflow/flow-go/storage"
    13  	"github.com/onflow/flow-go/storage/badger/transaction"
    14  	"github.com/onflow/flow-go/utils/unittest"
    15  )
    16  
    17  // TesKeyValueStoreStorage tests if the KV store is stored, retrieved and indexed correctly
    18  func TestKeyValueStoreStorage(t *testing.T) {
    19  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    20  		metrics := metrics.NewNoopCollector()
    21  		store := NewProtocolKVStore(metrics, db, DefaultProtocolKVStoreCacheSize, DefaultProtocolKVStoreByBlockIDCacheSize)
    22  
    23  		expected := &flow.PSKeyValueStoreData{
    24  			Version: 2,
    25  			Data:    unittest.RandomBytes(32),
    26  		}
    27  		stateID := unittest.IdentifierFixture()
    28  		blockID := unittest.IdentifierFixture()
    29  
    30  		// store protocol state and auxiliary info
    31  		err := transaction.Update(db, func(tx *transaction.Tx) error {
    32  			err := store.StoreTx(stateID, expected)(tx)
    33  			require.NoError(t, err)
    34  			return store.IndexTx(blockID, stateID)(tx)
    35  		})
    36  		require.NoError(t, err)
    37  
    38  		// fetch protocol state
    39  		actual, err := store.ByID(stateID)
    40  		require.NoError(t, err)
    41  		assert.Equal(t, expected, actual)
    42  
    43  		// fetch protocol state by block ID
    44  		actualByBlockID, err := store.ByBlockID(blockID)
    45  		require.NoError(t, err)
    46  		assert.Equal(t, expected, actualByBlockID)
    47  	})
    48  }
    49  
    50  // TestProtocolKVStore_StoreTx tests that StoreTx returns an error if the KV-store snapshot with the given id is already stored.
    51  func TestProtocolKVStore_StoreTx(t *testing.T) {
    52  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    53  		metrics := metrics.NewNoopCollector()
    54  		store := NewProtocolKVStore(metrics, db, DefaultProtocolKVStoreCacheSize, DefaultProtocolKVStoreByBlockIDCacheSize)
    55  
    56  		stateID := unittest.IdentifierFixture()
    57  		expected := &flow.PSKeyValueStoreData{
    58  			Version: 2,
    59  			Data:    unittest.RandomBytes(32),
    60  		}
    61  
    62  		err := transaction.Update(db, store.StoreTx(stateID, expected))
    63  		require.NoError(t, err)
    64  
    65  		err = transaction.Update(db, store.StoreTx(stateID, expected))
    66  		require.ErrorIs(t, err, storage.ErrAlreadyExists)
    67  	})
    68  }
    69  
    70  // TestProtocolKVStore_IndexTx tests that IndexTx returns an error if a KV store for the given blockID has already been indexed.
    71  func TestProtocolKVStore_IndexTx(t *testing.T) {
    72  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    73  		metrics := metrics.NewNoopCollector()
    74  		store := NewProtocolKVStore(metrics, db, DefaultProtocolKVStoreCacheSize, DefaultProtocolKVStoreByBlockIDCacheSize)
    75  
    76  		stateID := unittest.IdentifierFixture()
    77  		blockID := unittest.IdentifierFixture()
    78  
    79  		err := transaction.Update(db, store.IndexTx(blockID, stateID))
    80  		require.NoError(t, err)
    81  
    82  		err = transaction.Update(db, store.IndexTx(blockID, stateID))
    83  		require.ErrorIs(t, err, storage.ErrAlreadyExists)
    84  	})
    85  }
    86  
    87  // TestProtocolKVStore_ByBlockID tests that ByBlockID returns an error if no snapshot has been indexed for the given block.
    88  func TestProtocolKVStore_ByBlockID(t *testing.T) {
    89  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    90  		metrics := metrics.NewNoopCollector()
    91  		store := NewProtocolKVStore(metrics, db, DefaultProtocolKVStoreCacheSize, DefaultProtocolKVStoreByBlockIDCacheSize)
    92  
    93  		blockID := unittest.IdentifierFixture()
    94  		_, err := store.ByBlockID(blockID)
    95  		require.ErrorIs(t, err, storage.ErrNotFound)
    96  	})
    97  }
    98  
    99  // TestProtocolKVStore_ByID tests that ByID returns an error if no snapshot with the given Identifier is known.
   100  func TestProtocolKVStore_ByID(t *testing.T) {
   101  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
   102  		metrics := metrics.NewNoopCollector()
   103  		store := NewProtocolKVStore(metrics, db, DefaultProtocolKVStoreCacheSize, DefaultProtocolKVStoreByBlockIDCacheSize)
   104  
   105  		stateID := unittest.IdentifierFixture()
   106  		_, err := store.ByID(stateID)
   107  		require.ErrorIs(t, err, storage.ErrNotFound)
   108  	})
   109  }