From bb2357ec00a058eb628ae7a1b992d349452db2ba Mon Sep 17 00:00:00 2001
From: Dhruba Basu <7675102+dhrubabasu@users.noreply.github.com>
Date: Tue, 23 Jan 2024 18:21:03 -0500
Subject: [PATCH] Replace `closeFn` with `t.Cleanup` (#2638)

Co-authored-by: Stephen Buttolph <stephen@avalabs.org>
---
 database/rpcdb/db_test.go                     | 20 ++++++-------------
 .../gvalidators/validator_state_test.go       | 18 ++++++-----------
 vms/platformvm/warp/gwarp/signer_test.go      |  8 ++++----
 3 files changed, 16 insertions(+), 30 deletions(-)

diff --git a/database/rpcdb/db_test.go b/database/rpcdb/db_test.go
index a309ea01298..cc0cca3694b 100644
--- a/database/rpcdb/db_test.go
+++ b/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,8 +59,6 @@ func TestInterface(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			db := setupDB(t)
 			test(t, db.client)
-
-			db.closeFn()
 		})
 	}
 }
@@ -67,22 +66,16 @@ func TestInterface(t *testing.T) {
 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()
 			})
 		}
 	}
diff --git a/snow/validators/gvalidators/validator_state_test.go b/snow/validators/gvalidators/validator_state_test.go
index 9b6e692d864..0dbf9ebe895 100644
--- a/snow/validators/gvalidators/validator_state_test.go
+++ b/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()
diff --git a/vms/platformvm/warp/gwarp/signer_test.go b/vms/platformvm/warp/gwarp/signer_test.go
index 938d0ed07fa..31c7b3e993d 100644
--- a/vms/platformvm/warp/gwarp/signer_test.go
+++ b/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()
 		})
 	}
 }