minor test cleanup:

* moved `NewContext` outside of `WithReaderBatchWriter` in a few more cases, so that it's only released after the batch is committed.

Author: AlexHentschel

module/builder/collection/builder_test.go

@@ -80,7 +80,7 @@ type BuilderSuite struct {
 // runs before each test runs
 func (suite *BuilderSuite) SetupTest() {
 	fmt.Println("SetupTest>>>>")
-	lockManager := storage.NewTestingLockManager()
+	suite.lockManager = storage.NewTestingLockManager()
 	var err error
 
 	suite.genesis, err = unittest.ClusterBlock.Genesis()
@@ -92,7 +92,6 @@ func (suite *BuilderSuite) SetupTest() {
 	suite.dbdir = unittest.TempDir(suite.T())
 	suite.badgerDB = unittest.BadgerDB(suite.T(), suite.dbdir)
 	suite.db = badgerimpl.ToDB(suite.badgerDB)
-	suite.lockManager = lockManager
 
 	metrics := metrics.NewNoopCollector()
 	tracer := trace.NewNoopTracer()
@@ -132,16 +131,16 @@ func (suite *BuilderSuite) SetupTest() {
 	clusterQC := unittest.QuorumCertificateFixture(unittest.QCWithRootBlockID(suite.genesis.ID()))
 	clusterStateRoot, err := clusterkv.NewStateRoot(suite.genesis, clusterQC, suite.epochCounter)
 	suite.Require().NoError(err)
-	clusterState, err := clusterkv.Bootstrap(suite.db, lockManager, clusterStateRoot)
+	clusterState, err := clusterkv.Bootstrap(suite.db, suite.lockManager, clusterStateRoot)
 	suite.Require().NoError(err)
 
-	suite.state, err = clusterkv.NewMutableState(clusterState, lockManager, tracer, suite.headers, suite.payloads)
+	suite.state, err = clusterkv.NewMutableState(clusterState, suite.lockManager, tracer, suite.headers, suite.payloads)
 	suite.Require().NoError(err)
 
 	state, err := pbadger.Bootstrap(
 		metrics,
 		suite.db,
-		lockManager,
+		suite.lockManager,
 		all.Headers,
 		all.Seals,
 		all.Results,
@@ -178,7 +177,7 @@ func (suite *BuilderSuite) SetupTest() {
 		suite.Assert().True(added)
 	}
 
-	suite.builder, _ = builder.NewBuilder(suite.db, tracer, lockManager, suite.protoState, suite.state, suite.headers, suite.headers, suite.payloads, suite.pool, unittest.Logger(), suite.epochCounter)
+	suite.builder, _ = builder.NewBuilder(suite.db, tracer, suite.lockManager, suite.protoState, suite.state, suite.headers, suite.headers, suite.payloads, suite.pool, unittest.Logger(), suite.epochCounter)
 }
 
 // runs after each test finishes
@@ -201,17 +200,17 @@ func (suite *BuilderSuite) InsertBlock(block *model.Block) {
 }
 
 func (suite *BuilderSuite) FinalizeBlock(block model.Block) {
-	err := suite.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+	lctx := suite.lockManager.NewContext()
+	defer lctx.Release()
+	err := lctx.AcquireLock(storage.LockInsertOrFinalizeClusterBlock)
+	suite.Assert().NoError(err)
+
+	err = suite.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
 		var refBlock flow.Header
 		err := operation.RetrieveHeader(rw.GlobalReader(), block.Payload.ReferenceBlockID, &refBlock)
 		if err != nil {
 			return err
 		}
-		lctx := suite.lockManager.NewContext()
-		defer lctx.Release()
-		if err := lctx.AcquireLock(storage.LockInsertOrFinalizeClusterBlock); err != nil {
-			return err
-		}
 		err = procedure.FinalizeClusterBlock(lctx, rw, block.ID())
 		if err != nil {
 			return err
@@ -1208,7 +1207,6 @@ func benchmarkBuildOn(b *testing.B, size int) {
 		)
 		lctx := suite.lockManager.NewContext()
 		require.NoError(b, lctx.AcquireLock(storage.LockInsertOrFinalizeClusterBlock))
-
 		err := suite.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
 			return procedure.InsertClusterBlock(lctx, rw, unittest.ClusterProposalFromBlock(block))
 		})

module/builder/consensus/builder_test.go

@@ -178,7 +178,6 @@ func (bs *BuilderSuite) chainSeal(incorporatedResult *flow.IncorporatedResult) {
 // For the verifiers to start checking a result R, they need a source of randomness for the block _incorporating_
 // result R. The result for block [A3] is incorporated in [parent], which does _not_ have a child yet.
 func (bs *BuilderSuite) SetupTest() {
-
 	lockManager := storage.NewTestingLockManager()
 	// set up no-op dependencies
 	noopMetrics := metrics.NewNoopCollector()

module/finalizer/consensus/finalizer_test.go

@@ -64,14 +64,13 @@ func TestMakeFinalValidChain(t *testing.T) {
 	var list []flow.Identifier
 
 	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+		dbImpl := badgerimpl.ToDB(db)
+
 		// set up lock context
 		lockManager := storage.NewTestingLockManager()
 		lctx := lockManager.NewContext()
 		err := lctx.AcquireLock(storage.LockFinalizeBlock)
 		require.NoError(t, err)
-		defer lctx.Release()
-
-		dbImpl := badgerimpl.ToDB(db)
 
 		// insert the latest finalized height
 		err = dbImpl.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
@@ -84,9 +83,9 @@ func TestMakeFinalValidChain(t *testing.T) {
 			return operation.IndexFinalizedBlockByHeight(lctx, rw, final.Height, final.ID())
 		})
 		require.NoError(t, err)
+		lctx.Release()
 
 		// insert the finalized block header into the DB
-
 		insertLctx := lockManager.NewContext()
 		require.NoError(t, insertLctx.AcquireLock(storage.LockInsertBlock))
 		err = dbImpl.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
@@ -145,14 +144,13 @@ func TestMakeFinalInvalidHeight(t *testing.T) {
 	var list []flow.Identifier
 
 	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+		dbImpl := badgerimpl.ToDB(db)
+
 		// set up lock context
 		lockManager := storage.NewTestingLockManager()
 		lctx := lockManager.NewContext()
 		err := lctx.AcquireLock(storage.LockFinalizeBlock)
 		require.NoError(t, err)
-		defer lctx.Release()
-
-		dbImpl := badgerimpl.ToDB(db)
 
 		// insert the latest finalized height
 		err = dbImpl.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
@@ -165,6 +163,7 @@ func TestMakeFinalInvalidHeight(t *testing.T) {
 			return operation.IndexFinalizedBlockByHeight(lctx, rw, final.Height, final.ID())
 		})
 		require.NoError(t, err)
+		lctx.Release()
 
 		// insert the finalized block header into the DB
 		insertLctx := lockManager.NewContext()
@@ -219,14 +218,12 @@ func TestMakeFinalDuplicate(t *testing.T) {
 	var list []flow.Identifier
 
 	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+		dbImpl := badgerimpl.ToDB(db)
 		// set up lock context
 		lockManager := storage.NewTestingLockManager()
 		lctx := lockManager.NewContext()
 		err := lctx.AcquireLock(storage.LockFinalizeBlock)
 		require.NoError(t, err)
-		defer lctx.Release()
-
-		dbImpl := badgerimpl.ToDB(db)
 
 		// insert the latest finalized height
 		err = dbImpl.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
@@ -239,6 +236,7 @@ func TestMakeFinalDuplicate(t *testing.T) {
 			return operation.IndexFinalizedBlockByHeight(lctx, rw, final.Height, final.ID())
 		})
 		require.NoError(t, err)
+		lctx.Release()
 
 		// insert the finalized block header into the DB
 		insertLctx := lockManager.NewContext()

state/cluster/badger/mutator_test.go

@@ -66,8 +66,7 @@ func (suite *MutatorSuite) SetupTest() {
 	suite.dbdir = unittest.TempDir(suite.T())
 	suite.badgerdb = unittest.BadgerDB(suite.T(), suite.dbdir)
 	suite.db = badgerimpl.ToDB(suite.badgerdb)
-	lockManager := storage.NewTestingLockManager()
-	suite.lockManager = lockManager
+	suite.lockManager = storage.NewTestingLockManager()
 
 	metrics := metrics.NewNoopCollector()
 	tracer := trace.NewNoopTracer()
@@ -105,7 +104,7 @@ func (suite *MutatorSuite) SetupTest() {
 	state, err := pbadger.Bootstrap(
 		metrics,
 		suite.db,
-		lockManager,
+		suite.lockManager,
 		all.Headers,
 		all.Seals,
 		all.Results,

storage/store/commits_test.go

@@ -62,9 +62,9 @@ func TestCommitStoreAndRemove(t *testing.T) {
 		// Create and store a commit
 		blockID := unittest.IdentifierFixture()
 		expected := unittest.StateCommitmentFixture()
+		lctx := lockManager.NewContext()
+		defer lctx.Release()
 		err := db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-			lctx := lockManager.NewContext()
-			defer lctx.Release()
 			require.NoError(t, lctx.AcquireLock(storage.LockInsertOwnReceipt))
 			return store.BatchStore(lctx, blockID, expected, rw)
 		})