diff --git a/database/rpcdb/db_test.go b/database/rpcdb/db_test.go index a309ea01298b..cc0cca3694b0 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 9b6e692d8645..0dbf9ebe8952 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 938d0ed07fa5..31c7b3e993d4 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() }) } }