Use t.Cleanup in pebble.NewMemTest (#1278)

NethermindEth · Sep 27, 2023 · f0ab4e1 · f0ab4e1
1 parent c9e9b01
commit f0ab4e1
Show file tree

Hide file tree

Showing 16 changed files with 79 additions and 165 deletions.
diff --git a/blockchain/blockchain_test.go b/blockchain/blockchain_test.go
@@ -24,7 +24,7 @@ func TestNew(t *testing.T) {
 	gw := adaptfeeder.New(client)
 	log := utils.NewNopZapLogger()
 	t.Run("empty blockchain's head is nil", func(t *testing.T) {
-		chain := blockchain.New(pebble.NewMemTest(), utils.MAINNET, log)
+		chain := blockchain.New(pebble.NewMemTest(t), utils.MAINNET, log)
 		assert.Equal(t, utils.MAINNET, chain.Network())
 		b, err := chain.Head()
 		assert.Nil(t, b)
@@ -37,7 +37,7 @@ func TestNew(t *testing.T) {
 		stateUpdate0, err := gw.StateUpdate(context.Background(), 0)
 		require.NoError(t, err)
 
-		testDB := pebble.NewMemTest()
+		testDB := pebble.NewMemTest(t)
 		chain := blockchain.New(testDB, utils.MAINNET, log)
 		assert.NoError(t, chain.Store(block0, &emptyCommitments, stateUpdate0, nil))
 
@@ -53,7 +53,7 @@ func TestHeight(t *testing.T) {
 	gw := adaptfeeder.New(client)
 	log := utils.NewNopZapLogger()
 	t.Run("return nil if blockchain is empty", func(t *testing.T) {
-		chain := blockchain.New(pebble.NewMemTest(), utils.GOERLI, log)
+		chain := blockchain.New(pebble.NewMemTest(t), utils.GOERLI, log)
 		_, err := chain.Height()
 		assert.Error(t, err)
 	})
@@ -64,7 +64,7 @@ func TestHeight(t *testing.T) {
 		stateUpdate0, err := gw.StateUpdate(context.Background(), 0)
 		require.NoError(t, err)
 
-		testDB := pebble.NewMemTest()
+		testDB := pebble.NewMemTest(t)
 		chain := blockchain.New(testDB, utils.MAINNET, log)
 		assert.NoError(t, chain.Store(block0, &emptyCommitments, stateUpdate0, nil))
 
@@ -76,7 +76,7 @@ func TestHeight(t *testing.T) {
 }
 
 func TestBlockByNumberAndHash(t *testing.T) {
-	chain := blockchain.New(pebble.NewMemTest(), utils.GOERLI, utils.NewNopZapLogger())
+	chain := blockchain.New(pebble.NewMemTest(t), utils.GOERLI, utils.NewNopZapLogger())
 	t.Run("same block is returned for both GetBlockByNumber and GetBlockByHash", func(t *testing.T) {
 		client := feeder.NewTestClient(t, utils.MAINNET)
 		gw := adaptfeeder.New(client)
@@ -112,7 +112,7 @@ func TestVerifyBlock(t *testing.T) {
 	h1, err := new(felt.Felt).SetRandom()
 	require.NoError(t, err)
 
-	chain := blockchain.New(pebble.NewMemTest(), utils.MAINNET, utils.NewNopZapLogger())
+	chain := blockchain.New(pebble.NewMemTest(t), utils.MAINNET, utils.NewNopZapLogger())
 
 	t.Run("error if chain is empty and incoming block number is not 0", func(t *testing.T) {
 		block := &core.Block{Header: &core.Header{Number: 10}}
@@ -174,7 +174,7 @@ func TestSanityCheckNewHeight(t *testing.T) {
 	h1, err := new(felt.Felt).SetRandom()
 	require.NoError(t, err)
 
-	chain := blockchain.New(pebble.NewMemTest(), utils.MAINNET, utils.NewNopZapLogger())
+	chain := blockchain.New(pebble.NewMemTest(t), utils.MAINNET, utils.NewNopZapLogger())
 
 	client := feeder.NewTestClient(t, utils.MAINNET)
 
@@ -220,7 +220,7 @@ func TestStore(t *testing.T) {
 	require.NoError(t, err)
 
 	t.Run("add block to empty blockchain", func(t *testing.T) {
-		chain := blockchain.New(pebble.NewMemTest(), utils.MAINNET, log)
+		chain := blockchain.New(pebble.NewMemTest(t), utils.MAINNET, log)
 		require.NoError(t, chain.Store(block0, &emptyCommitments, stateUpdate0, nil))
 
 		headBlock, err := chain.Head()
@@ -246,7 +246,7 @@ func TestStore(t *testing.T) {
 		stateUpdate1, err := gw.StateUpdate(context.Background(), 1)
 		require.NoError(t, err)
 
-		chain := blockchain.New(pebble.NewMemTest(), utils.MAINNET, log)
+		chain := blockchain.New(pebble.NewMemTest(t), utils.MAINNET, log)
 		require.NoError(t, chain.Store(block0, &emptyCommitments, stateUpdate0, nil))
 		require.NoError(t, chain.Store(block1, &emptyCommitments, stateUpdate1, nil))
 
@@ -269,7 +269,7 @@ func TestStore(t *testing.T) {
 }
 
 func TestTransactionAndReceipt(t *testing.T) {
-	chain := blockchain.New(pebble.NewMemTest(), utils.MAINNET, utils.NewNopZapLogger())
+	chain := blockchain.New(pebble.NewMemTest(t), utils.MAINNET, utils.NewNopZapLogger())
 
 	client := feeder.NewTestClient(t, utils.MAINNET)
 	gw := adaptfeeder.New(client)
@@ -356,10 +356,7 @@ func TestTransactionAndReceipt(t *testing.T) {
 }
 
 func TestState(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 	chain := blockchain.New(testDB, utils.MAINNET, utils.NewNopZapLogger())
 
 	client := feeder.NewTestClient(t, utils.MAINNET)
@@ -422,10 +419,7 @@ func TestState(t *testing.T) {
 }
 
 func TestEvents(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 	chain := blockchain.New(testDB, utils.GOERLI2, utils.NewNopZapLogger())
 
 	client := feeder.NewTestClient(t, utils.GOERLI2)
@@ -544,7 +538,7 @@ func TestEvents(t *testing.T) {
 }
 
 func TestRevert(t *testing.T) {
-	testdb := pebble.NewMemTest()
+	testdb := pebble.NewMemTest(t)
 	chain := blockchain.New(testdb, utils.MAINNET, utils.NewNopZapLogger())
 
 	client := feeder.NewTestClient(t, utils.MAINNET)
@@ -606,7 +600,7 @@ func TestRevert(t *testing.T) {
 				return err
 			}
 			assert.False(t, it.Next(), it.Key())
-			return nil
+			return it.Close()
 		}))
 	})
 
@@ -629,7 +623,7 @@ func TestL1Update(t *testing.T) {
 
 	for _, head := range heads {
 		t.Run(fmt.Sprintf("update L1 head to block %d", head.BlockNumber), func(t *testing.T) {
-			chain := blockchain.New(pebble.NewMemTest(), utils.MAINNET, utils.NewNopZapLogger())
+			chain := blockchain.New(pebble.NewMemTest(t), utils.MAINNET, utils.NewNopZapLogger())
 			require.NoError(t, chain.SetL1Head(head))
 			got, err := chain.L1Head()
 			require.NoError(t, err)
@@ -639,10 +633,7 @@ func TestL1Update(t *testing.T) {
 }
 
 func TestPending(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 	chain := blockchain.New(testDB, utils.MAINNET, utils.NewNopZapLogger())
 	client := feeder.NewTestClient(t, utils.MAINNET)
 	gw := adaptfeeder.New(client)
@@ -769,10 +760,7 @@ func TestPending(t *testing.T) {
 }
 
 func TestSubscribeNewHeads(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 	chain := blockchain.New(testDB, utils.MAINNET, utils.NewNopZapLogger())
 	client := feeder.NewTestClient(t, utils.MAINNET)
 	gw := adaptfeeder.New(client)

diff --git a/core/contract_test.go b/core/contract_test.go
@@ -52,10 +52,7 @@ func TestContractAddress(t *testing.T) {
 }
 
 func TestNewContract(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
@@ -129,10 +126,7 @@ func TestNewContract(t *testing.T) {
 }
 
 func TestNonceAndClassHash(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 
 	txn := testDB.NewTransaction(true)
 	addr := new(felt.Felt).SetUint64(44)
@@ -170,10 +164,7 @@ func TestNonceAndClassHash(t *testing.T) {
 }
 
 func TestUpdateStorageAndStorage(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 
 	txn := testDB.NewTransaction(true)
 	addr := new(felt.Felt).SetUint64(44)
@@ -211,10 +202,7 @@ func TestUpdateStorageAndStorage(t *testing.T) {
 }
 
 func TestPurge(t *testing.T) {
-	testDB := pebble.NewMemTest()
-	t.Cleanup(func() {
-		require.NoError(t, testDB.Close())
-	})
+	testDB := pebble.NewMemTest(t)
 
 	txn := testDB.NewTransaction(true)
 	addr := new(felt.Felt).SetUint64(44)

diff --git a/core/history_pkg_test.go b/core/history_pkg_test.go
@@ -10,11 +10,10 @@ import (
 )
 
 func TestHistory(t *testing.T) {
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
-		require.NoError(t, testDB.Close())
 	})
 
 	history := &history{txn: txn}

diff --git a/core/state_test.go b/core/state_test.go
@@ -23,7 +23,7 @@ func TestUpdate(t *testing.T) {
 	client := feeder.NewTestClient(t, utils.MAINNET)
 	gw := adaptfeeder.New(client)
 
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -152,7 +152,7 @@ func TestContractClassHash(t *testing.T) {
 	client := feeder.NewTestClient(t, utils.MAINNET)
 	gw := adaptfeeder.New(client)
 
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -211,7 +211,7 @@ func TestContractClassHash(t *testing.T) {
 }
 
 func TestNonce(t *testing.T) {
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -262,7 +262,7 @@ func TestNonce(t *testing.T) {
 }
 
 func TestStateHistory(t *testing.T) {
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -316,7 +316,7 @@ func TestContractIsDeployedAt(t *testing.T) {
 	client := feeder.NewTestClient(t, utils.MAINNET)
 	gw := adaptfeeder.New(client)
 
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -364,7 +364,7 @@ func TestContractIsDeployedAt(t *testing.T) {
 }
 
 func TestClass(t *testing.T) {
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -408,7 +408,7 @@ func TestClass(t *testing.T) {
 }
 
 func TestRevert(t *testing.T) {
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -555,7 +555,7 @@ func TestRevert(t *testing.T) {
 				return err
 			}
 			assert.False(t, it.Next())
-			return nil
+			return it.Close()
 		}))
 	})
 }
@@ -564,7 +564,7 @@ func TestRevertNoClassContracts(t *testing.T) {
 	client := feeder.NewTestClient(t, utils.MAINNET)
 	gw := adaptfeeder.New(client)
 
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())
@@ -602,7 +602,7 @@ func TestRevertNoClassContracts(t *testing.T) {
 }
 
 func TestRevertDeclaredClasses(t *testing.T) {
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	txn := testDB.NewTransaction(true)
 	t.Cleanup(func() {
 		require.NoError(t, txn.Discard())

diff --git a/core/trie/transaction_storage_test.go b/core/trie/transaction_storage_test.go
@@ -13,7 +13,7 @@ import (
 )
 
 func TestTransactionStorage(t *testing.T) {
-	testDB := pebble.NewMemTest()
+	testDB := pebble.NewMemTest(t)
 	prefix := []byte{37, 44}
 	key := trie.NewKey(44, nil)
 

diff --git a/db/pebble/db.go b/db/pebble/db.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"sync"
+	"testing"
 
 	"github.com/NethermindEth/juno/db"
 	"github.com/NethermindEth/juno/utils"
@@ -38,11 +39,16 @@ func NewMem() (db.DB, error) {
 }
 
 // NewMemTest opens a new in-memory database, panics on error
-func NewMemTest() db.DB {
+func NewMemTest(t *testing.T) db.DB {
 	memDB, err := NewMem()
 	if err != nil {
-		panic(err)
+		t.Fatalf("create in-memory db: %v", err)
 	}
+	t.Cleanup(func() {
+		if err := memDB.Close(); err != nil {
+			t.Errorf("close in-memory db: %v", err)
+		}
+	})
 	return memDB
 }