Replace closeFn with t.Cleanup (#2638)

Co-authored-by: Stephen Buttolph <[email protected]>

database/rpcdb/db_test.go

@@ -19,9 +19,8 @@ import (
 )
 
 type testDatabase struct {
-	client  *DatabaseClient
-	server  *memdb.Database
-	closeFn func()
+	client *DatabaseClient
+	server *memdb.Database
 }
 
 func setupDB(t testing.TB) *testDatabase {
@@ -45,11 +44,13 @@ func setupDB(t testing.TB) *testDatabase {
 	require.NoError(err)
 
 	db.client = NewClient(rpcdbpb.NewDatabaseClient(conn))
-	db.closeFn = func() {
+
+	t.Cleanup(func() {
 		serverCloser.Stop()
 		_ = conn.Close()
 		_ = listener.Close()
-	}
+	})
+
 	return db
 }
 
@@ -58,31 +59,23 @@ func TestInterface(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			db := setupDB(t)
 			test(t, db.client)
-
-			db.closeFn()
 		})
 	}
 }
 
 func FuzzKeyValue(f *testing.F) {
 	db := setupDB(f)
 	database.FuzzKeyValue(f, db.client)
-
-	db.closeFn()
 }
 
 func FuzzNewIteratorWithPrefix(f *testing.F) {
 	db := setupDB(f)
 	database.FuzzNewIteratorWithPrefix(f, db.client)
-
-	db.closeFn()
 }
 
 func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
 	db := setupDB(f)
 	database.FuzzNewIteratorWithStartAndPrefix(f, db.client)
-
-	db.closeFn()
 }
 
 func BenchmarkInterface(b *testing.B) {
@@ -92,7 +85,6 @@ func BenchmarkInterface(b *testing.B) {
 			b.Run(fmt.Sprintf("rpcdb_%d_pairs_%d_keys_%d_values_%s", size[0], size[1], size[2], name), func(b *testing.B) {
 				db := setupDB(b)
 				bench(b, db.client, keys, values)
-				db.closeFn()
 			})
 		}
 	}

snow/validators/gvalidators/validator_state_test.go

@@ -24,9 +24,8 @@ import (
 var errCustom = errors.New("custom")
 
 type testState struct {
-	client  *Client
-	server  *validators.MockState
-	closeFn func()
+	client *Client
+	server *validators.MockState
 }
 
 func setupState(t testing.TB, ctrl *gomock.Controller) *testState {
@@ -52,11 +51,13 @@ func setupState(t testing.TB, ctrl *gomock.Controller) *testState {
 	require.NoError(err)
 
 	state.client = NewClient(pb.NewValidatorStateClient(conn))
-	state.closeFn = func() {
+
+	t.Cleanup(func() {
 		serverCloser.Stop()
 		_ = conn.Close()
 		_ = listener.Close()
-	}
+	})
+
 	return state
 }
 
@@ -65,7 +66,6 @@ func TestGetMinimumHeight(t *testing.T) {
 	ctrl := gomock.NewController(t)
 
 	state := setupState(t, ctrl)
-	defer state.closeFn()
 
 	// Happy path
 	expectedHeight := uint64(1337)
@@ -88,7 +88,6 @@ func TestGetCurrentHeight(t *testing.T) {
 	ctrl := gomock.NewController(t)
 
 	state := setupState(t, ctrl)
-	defer state.closeFn()
 
 	// Happy path
 	expectedHeight := uint64(1337)
@@ -111,7 +110,6 @@ func TestGetSubnetID(t *testing.T) {
 	ctrl := gomock.NewController(t)
 
 	state := setupState(t, ctrl)
-	defer state.closeFn()
 
 	// Happy path
 	chainID := ids.GenerateTestID()
@@ -135,7 +133,6 @@ func TestGetValidatorSet(t *testing.T) {
 	ctrl := gomock.NewController(t)
 
 	state := setupState(t, ctrl)
-	defer state.closeFn()
 
 	// Happy path
 	sk0, err := bls.NewSecretKey()
@@ -209,9 +206,6 @@ func benchmarkGetValidatorSet(b *testing.B, vs map[ids.NodeID]*validators.GetVal
 	require := require.New(b)
 	ctrl := gomock.NewController(b)
 	state := setupState(b, ctrl)
-	defer func() {
-		state.closeFn()
-	}()
 
 	height := uint64(1337)
 	subnetID := ids.GenerateTestID()

vms/platformvm/warp/gwarp/signer_test.go

@@ -23,7 +23,6 @@ type testSigner struct {
 	sk        *bls.SecretKey
 	networkID uint32
 	chainID   ids.ID
-	closeFn   func()
 }
 
 func setupSigner(t testing.TB) *testSigner {
@@ -55,11 +54,13 @@ func setupSigner(t testing.TB) *testSigner {
 	require.NoError(err)
 
 	s.client = NewClient(pb.NewSignerClient(conn))
-	s.closeFn = func() {
+
+	t.Cleanup(func() {
 		serverCloser.Stop()
 		_ = conn.Close()
 		_ = listener.Close()
-	}
+	})
+
 	return s
 }
 
@@ -68,7 +69,6 @@ func TestInterface(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			s := setupSigner(t)
 			test(t, s.client, s.sk, s.networkID, s.chainID)
-			s.closeFn()
 		})
 	}
 }