Ensure DB connections are closed in tests (#6415)

## Motivation

All tests have been updated to ensure that DB connections are closed at the end of the test. This ensures that any open connections that weren't returned to the pool at the end of the test cause it to fail and would otherwise go unnoticed.

activation/activation_test.go

@@ -77,8 +77,8 @@ func newTestBuilder(tb testing.TB, numSigners int, opts ...BuilderOption) *testA
 
 	ctrl := gomock.NewController(tb)
 	tab := &testAtxBuilder{
-		db:          statesql.InMemory(),
-		localDb:     localsql.InMemory(sql.WithConnections(numSigners)),
+		db:          statesql.InMemoryTest(tb),
+		localDb:     localsql.InMemoryTest(tb, sql.WithConnections(numSigners)),
 		goldenATXID: types.ATXID(types.HexToHash32("77777")),
 
 		observedLogs: observedLogs,

activation/certifier_test.go

@@ -22,7 +22,7 @@ import (
 func TestPersistsCerts(t *testing.T) {
 	client := NewMockcertifierClient(gomock.NewController(t))
 	id := types.RandomNodeID()
-	db := localsql.InMemory()
+	db := localsql.InMemoryTest(t)
 	cert := &certdb.PoetCert{Data: []byte("cert"), Signature: []byte("sig")}
 	certifierAddress := &url.URL{Scheme: "http", Host: "certifier.org"}
 	pubkey := []byte("pubkey")
@@ -56,7 +56,7 @@ func TestPersistsCerts(t *testing.T) {
 }
 
 func TestAvoidsRedundantQueries(t *testing.T) {
-	db := localsql.InMemory()
+	db := localsql.InMemoryTest(t)
 	client := NewMockcertifierClient(gomock.NewController(t))
 	id1 := types.RandomNodeID()
 	id2 := types.RandomNodeID()
@@ -114,16 +114,16 @@ func TestObtainingPost(t *testing.T) {
 	id := types.RandomNodeID()
 
 	t.Run("no POST or ATX", func(t *testing.T) {
-		db := statesql.InMemory()
-		localDb := localsql.InMemory()
+		db := statesql.InMemoryTest(t)
+		localDb := localsql.InMemoryTest(t)
 
 		certifier := NewCertifierClient(db, localDb, zaptest.NewLogger(t))
 		_, err := certifier.obtainPost(context.Background(), id)
 		require.ErrorContains(t, err, "PoST not found")
 	})
 	t.Run("initial POST available", func(t *testing.T) {
-		db := statesql.InMemory()
-		localDb := localsql.InMemory()
+		db := statesql.InMemoryTest(t)
+		localDb := localsql.InMemoryTest(t)
 
 		post := nipost.Post{
 			Nonce:         30,
@@ -143,8 +143,8 @@ func TestObtainingPost(t *testing.T) {
 		require.Equal(t, post, *got)
 	})
 	t.Run("initial POST unavailable but ATX exists", func(t *testing.T) {
-		db := statesql.InMemory()
-		localDb := localsql.InMemory()
+		db := statesql.InMemoryTest(t)
+		localDb := localsql.InMemoryTest(t)
 
 		atx := newInitialATXv1(t, types.RandomATXID())
 		atx.SmesherID = id

activation/e2e/activation_test.go

@@ -33,8 +33,8 @@ import (
 	"github.com/spacemeshos/go-spacemesh/timesync"
 )
 
-func syncedSyncer(t testing.TB) *activation.Mocksyncer {
-	syncer := activation.NewMocksyncer(gomock.NewController(t))
+func syncedSyncer(tb testing.TB) *activation.Mocksyncer {
+	syncer := activation.NewMocksyncer(gomock.NewController(tb))
 	syncer.EXPECT().RegisterForATXSynced().DoAndReturn(func() <-chan struct{} {
 		synced := make(chan struct{})
 		close(synced)
@@ -61,8 +61,8 @@ func Test_BuilderWithMultipleClients(t *testing.T) {
 	logger := zaptest.NewLogger(t)
 	goldenATX := types.ATXID{2, 3, 4}
 	cfg := testPostConfig()
-	db := statesql.InMemory()
-	localDB := localsql.InMemory()
+	db := statesql.InMemoryTest(t)
+	localDB := localsql.InMemoryTest(t)
 
 	svc := grpcserver.NewPostService(logger, grpcserver.PostServiceQueryInterval(100*time.Millisecond))
 	svc.AllowConnections(true)

activation/e2e/atx_merge_test.go

@@ -41,16 +41,16 @@ const (
 	testTickSize = 1
 )
 
-func constructMerkleProof(t testing.TB, members []types.Hash32, ids map[uint64]bool) wire.MerkleProofV2 {
-	t.Helper()
+func constructMerkleProof(tb testing.TB, members []types.Hash32, ids map[uint64]bool) wire.MerkleProofV2 {
+	tb.Helper()
 
 	tree, err := merkle.NewTreeBuilder().
 		WithLeavesToProve(ids).
 		WithHashFunc(shared.HashMembershipTreeNode).
 		Build()
-	require.NoError(t, err)
+	require.NoError(tb, err)
 	for _, member := range members {
-		require.NoError(t, tree.AddLeaf(member[:]))
+		require.NoError(tb, tree.AddLeaf(member[:]))
 	}
 	nodes := tree.Proof()
 	nodesH32 := make([]types.Hash32, 0, len(nodes))
@@ -140,7 +140,7 @@ func createSoloAtx(publish types.EpochID, prev, pos types.ATXID, nipost *nipost.
 }
 
 func createMerged(
-	t testing.TB,
+	tb testing.TB,
 	niposts []nipostData,
 	publish types.EpochID,
 	marriage, positioning types.ATXID,
@@ -162,7 +162,7 @@ func createMerged(
 	// Append PoSTs for all IDs
 	for i, nipost := range niposts {
 		idx := slices.IndexFunc(previous, func(a types.ATXID) bool { return a == nipost.previous })
-		require.NotEqual(t, -1, idx)
+		require.NotEqual(tb, -1, idx)
 		atx.NiPosts[0].Posts = append(atx.NiPosts[0].Posts, wire.SubPostV2{
 			MarriageIndex:       uint32(i),
 			PrevATXIndex:        uint32(idx),
@@ -174,16 +174,16 @@ func createMerged(
 	return atx
 }
 
-func signers(t testing.TB, keysHex []string) []*signing.EdSigner {
-	t.Helper()
+func signers(tb testing.TB, keysHex []string) []*signing.EdSigner {
+	tb.Helper()
 
 	signers := make([]*signing.EdSigner, 0, len(keysHex))
 	for _, k := range keysHex {
 		key, err := hex.DecodeString(k)
-		require.NoError(t, err)
+		require.NoError(tb, err)
 
 		sig, err := signing.NewEdSigner(signing.WithPrivateKey(key))
-		require.NoError(t, err)
+		require.NoError(tb, err)
 		signers = append(signers, sig)
 	}
 	return signers
@@ -211,6 +211,7 @@ func Test_MarryAndMerge(t *testing.T) {
 	cfg := testPostConfig()
 	db := statesql.InMemoryTest(t)
 	cdb := datastore.NewCachedDB(db, logger)
+	t.Cleanup(func() { assert.NoError(t, cdb.Close()) })
 	localDB := localsql.InMemoryTest(t)
 
 	svc := grpcserver.NewPostService(logger, grpcserver.PostServiceQueryInterval(100*time.Millisecond))

activation/e2e/builds_atx_v2_test.go

@@ -53,8 +53,9 @@ func TestBuilder_SwitchesToBuildV2(t *testing.T) {
 	require.NoError(t, err)
 
 	cfg := testPostConfig()
-	db := statesql.InMemory()
+	db := statesql.InMemoryTest(t)
 	cdb := datastore.NewCachedDB(db, logger)
+	t.Cleanup(func() { assert.NoError(t, cdb.Close()) })
 
 	opts := testPostSetupOpts(t)
 	svc := grpcserver.NewPostService(logger, grpcserver.PostServiceQueryInterval(100*time.Millisecond))
@@ -96,7 +97,7 @@ func TestBuilder_SwitchesToBuildV2(t *testing.T) {
 	client := ae2e.NewTestPoetClient(1, poetCfg)
 	poetClient := activation.NewPoetServiceWithClient(poetDb, client, poetCfg, logger, testTickSize)
 
-	localDB := localsql.InMemory()
+	localDB := localsql.InMemoryTest(t)
 	nb, err := activation.NewNIPostBuilder(
 		localDB,
 		svc,

activation/e2e/certifier_client_test.go

@@ -34,8 +34,8 @@ func TestCertification(t *testing.T) {
 	require.NoError(t, err)
 
 	cfg := testPostConfig()
-	db := statesql.InMemory()
-	localDb := localsql.InMemory()
+	db := statesql.InMemoryTest(t)
+	localDb := localsql.InMemoryTest(t)
 
 	opts := testPostSetupOpts(t)
 	logger := zaptest.NewLogger(t)
@@ -171,26 +171,26 @@ func (c *testCertifier) certify(w http.ResponseWriter, r *http.Request) {
 }
 
 func spawnTestCertifier(
-	t *testing.T,
+	tb testing.TB,
 	cfg activation.PostConfig,
 	// optional - if nil, will create valid certs
 	makeCert func(nodeID []byte) *poetShared.Cert,
 	opts ...verifying.OptionFunc,
 ) (ed25519.PublicKey, net.Addr) {
-	t.Helper()
+	tb.Helper()
 
 	pub, private, err := ed25519.GenerateKey(nil)
-	require.NoError(t, err)
+	require.NoError(tb, err)
 
 	postVerifier, err := activation.NewPostVerifier(
 		cfg,
-		zaptest.NewLogger(t),
+		zaptest.NewLogger(tb),
 		activation.WithVerifyingOpts(activation.DefaultTestPostVerifyingOpts()),
 	)
-	require.NoError(t, err)
+	require.NoError(tb, err)
 	var eg errgroup.Group
 	l, err := net.Listen("tcp", ":0")
-	require.NoError(t, err)
+	require.NoError(tb, err)
 	eg.Go(func() error {
 		certifier := &testCertifier{
 			privKey:      private,
@@ -205,9 +205,9 @@ func spawnTestCertifier(
 		http.Serve(l, mux)
 		return nil
 	})
-	t.Cleanup(func() {
-		require.NoError(t, l.Close())
-		require.NoError(t, eg.Wait())
+	tb.Cleanup(func() {
+		require.NoError(tb, l.Close())
+		require.NoError(tb, eg.Wait())
 	})
 
 	return pub, l.Addr()

activation/e2e/checkpoint_merged_test.go

@@ -45,6 +45,7 @@ func Test_CheckpointAfterMerge(t *testing.T) {
 	cfg := testPostConfig()
 	db := statesql.InMemoryTest(t)
 	cdb := datastore.NewCachedDB(db, logger)
+	t.Cleanup(func() { assert.NoError(t, cdb.Close()) })
 	localDB := localsql.InMemoryTest(t)
 
 	svc := grpcserver.NewPostService(logger, grpcserver.PostServiceQueryInterval(100*time.Millisecond))
@@ -283,14 +284,15 @@ func Test_CheckpointAfterMerge(t *testing.T) {
 	// 4. Spawn new ATX handler and builder using the new DB
 	poetDb, err = activation.NewPoetDb(newDB, logger.Named("poetDb"))
 	require.NoError(t, err)
-	cdb = datastore.NewCachedDB(newDB, logger)
+	newCdb := datastore.NewCachedDB(newDB, logger)
+	t.Cleanup(func() { assert.NoError(t, newCdb.Close()) })
 
 	poetSvc = activation.NewPoetServiceWithClient(poetDb, client, poetCfg, logger, testTickSize)
 	validator = activation.NewValidator(newDB, poetDb, cfg, opts.Scrypt, verifier)
 	require.NoError(t, err)
 	atxHdlr = activation.NewHandler(
 		"local",
-		cdb,
+		newCdb,
 		atxsdata.New(),
 		signing.NewEdVerifier(),
 		clock,

activation/e2e/checkpoint_test.go

@@ -48,6 +48,7 @@ func TestCheckpoint_PublishingSoloATXs(t *testing.T) {
 	cfg := testPostConfig()
 	db := statesql.InMemoryTest(t)
 	cdb := datastore.NewCachedDB(db, logger)
+	t.Cleanup(func() { assert.NoError(t, cdb.Close()) })
 
 	opts := testPostSetupOpts(t)
 	svc := grpcserver.NewPostService(logger, grpcserver.PostServiceQueryInterval(100*time.Millisecond))
@@ -189,14 +190,15 @@ func TestCheckpoint_PublishingSoloATXs(t *testing.T) {
 	// 3. Spawn new ATX handler and builder using the new DB
 	poetDb, err = activation.NewPoetDb(newDB, logger.Named("poetDb"))
 	require.NoError(t, err)
-	cdb = datastore.NewCachedDB(newDB, logger)
+	newCdb := datastore.NewCachedDB(newDB, logger)
+	t.Cleanup(func() { assert.NoError(t, newCdb.Close()) })
 	atxdata, err = atxsdata.Warm(newDB, 1, logger, sig)
 	poetService = activation.NewPoetServiceWithClient(poetDb, client, poetCfg, logger, testTickSize)
 	validator = activation.NewValidator(newDB, poetDb, cfg, opts.Scrypt, verifier)
 	require.NoError(t, err)
 	atxHdlr = activation.NewHandler(
 		"local",
-		cdb,
+		newCdb,
 		atxdata,
 		edVerifier,
 		clock,

activation/e2e/nipost_test.go

@@ -136,7 +136,7 @@ func initPost(
 
 	logger := zaptest.NewLogger(tb)
 	syncer := syncedSyncer(tb)
-	db := statesql.InMemory()
+	db := statesql.InMemoryTest(tb)
 	mgr, err := activation.NewPostSetupManager(cfg, logger, db, atxsdata.New(), golden, syncer, nil)
 	require.NoError(tb, err)
 
@@ -157,8 +157,8 @@ func TestNIPostBuilderWithClients(t *testing.T) {
 	logger := zaptest.NewLogger(t)
 	goldenATX := types.ATXID{2, 3, 4}
 	cfg := testPostConfig()
-	db := statesql.InMemory()
-	localDb := localsql.InMemory()
+	db := statesql.InMemoryTest(t)
+	localDb := localsql.InMemoryTest(t)
 
 	opts := testPostSetupOpts(t)
 	svc := grpcserver.NewPostService(logger, grpcserver.PostServiceQueryInterval(100*time.Millisecond))
@@ -202,7 +202,7 @@ func TestNIPostBuilderWithClients(t *testing.T) {
 	client := ae2e.NewTestPoetClient(1, poetCfg)
 	poetService := activation.NewPoetServiceWithClient(poetDb, client, poetCfg, logger, testTickSize)
 
-	localDB := localsql.InMemory()
+	localDB := localsql.InMemoryTest(t)
 	nb, err := activation.NewNIPostBuilder(
 		localDB,
 		svc,
@@ -244,7 +244,7 @@ func Test_NIPostBuilderWithMultipleClients(t *testing.T) {
 	logger := zaptest.NewLogger(t)
 	goldenATX := types.ATXID{2, 3, 4}
 	cfg := testPostConfig()
-	db := statesql.InMemory()
+	db := statesql.InMemoryTest(t)
 
 	opts := testPostSetupOpts(t)
 	svc := grpcserver.NewPostService(logger, grpcserver.PostServiceQueryInterval(100*time.Millisecond))
@@ -288,7 +288,7 @@ func Test_NIPostBuilderWithMultipleClients(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { assert.NoError(t, verifier.Close()) })
 
-	localDB := localsql.InMemory()
+	localDB := localsql.InMemoryTest(t)
 	nb, err := activation.NewNIPostBuilder(
 		localDB,
 		svc,

activation/e2e/validation_test.go

@@ -29,7 +29,7 @@ func TestValidator_Validate(t *testing.T) {
 	logger := zaptest.NewLogger(t)
 	goldenATX := types.ATXID{2, 3, 4}
 	cfg := testPostConfig()
-	db := statesql.InMemory()
+	db := statesql.InMemoryTest(t)
 
 	validator := activation.NewMocknipostValidator(gomock.NewController(t))
 
@@ -50,7 +50,7 @@ func TestValidator_Validate(t *testing.T) {
 		GracePeriod: epoch / 4,
 	}
 
-	poetDb, err := activation.NewPoetDb(statesql.InMemory(), logger.Named("poetDb"))
+	poetDb, err := activation.NewPoetDb(db, logger.Named("poetDb"))
 	require.NoError(t, err)
 	client := ae2e.NewTestPoetClient(1, poetCfg)
 	poetService := activation.NewPoetServiceWithClient(poetDb, client, poetCfg, logger, testTickSize)
@@ -68,7 +68,7 @@ func TestValidator_Validate(t *testing.T) {
 
 	challenge := types.RandomHash()
 	nb, err := activation.NewNIPostBuilder(
-		localsql.InMemory(),
+		localsql.InMemoryTest(t),
 		svc,
 		logger.Named("nipostBuilder"),
 		poetCfg,