From 2a491ec3f054155d1e9e1755a6aa4925a61dce38 Mon Sep 17 00:00:00 2001
From: Joachim Bauch <bauch@struktur.de>
Date: Tue, 5 Apr 2022 12:57:20 +0200
Subject: [PATCH] Use "T.Cleanup()" instead of manually deferred shutdown
 function.

---
 backend_server_test.go      |  57 +++++++----------
 clientsession_test.go       |   6 +-
 hub_test.go                 | 121 +++++++++++++-----------------------
 natsclient_loopback_test.go |  19 +++---
 natsclient_test.go          |  30 ++++-----
 proxy/proxy_server_test.go  |  11 ++--
 room_test.go                |  12 ++--
 transient_data_test.go      |   3 +-
 virtualsession_test.go      |   6 +-
 9 files changed, 99 insertions(+), 166 deletions(-)

diff --git a/backend_server_test.go b/backend_server_test.go
index cd6dba2..82716dd 100644
--- a/backend_server_test.go
+++ b/backend_server_test.go
@@ -52,11 +52,11 @@ var (
 	turnServers       = strings.Split(turnServersString, ",")
 )
 
-func CreateBackendServerForTest(t *testing.T) (*goconf.ConfigFile, *BackendServer, NatsClient, *Hub, *mux.Router, *httptest.Server, func()) {
+func CreateBackendServerForTest(t *testing.T) (*goconf.ConfigFile, *BackendServer, NatsClient, *Hub, *mux.Router, *httptest.Server) {
 	return CreateBackendServerForTestFromConfig(t, nil)
 }
 
-func CreateBackendServerForTestWithTurn(t *testing.T) (*goconf.ConfigFile, *BackendServer, NatsClient, *Hub, *mux.Router, *httptest.Server, func()) {
+func CreateBackendServerForTestWithTurn(t *testing.T) (*goconf.ConfigFile, *BackendServer, NatsClient, *Hub, *mux.Router, *httptest.Server) {
 	config := goconf.NewConfigFile()
 	config.AddOption("turn", "apikey", turnApiKey)
 	config.AddOption("turn", "secret", turnSecret)
@@ -64,7 +64,7 @@ func CreateBackendServerForTestWithTurn(t *testing.T) (*goconf.ConfigFile, *Back
 	return CreateBackendServerForTestFromConfig(t, config)
 }
 
-func CreateBackendServerForTestFromConfig(t *testing.T, config *goconf.ConfigFile) (*goconf.ConfigFile, *BackendServer, NatsClient, *Hub, *mux.Router, *httptest.Server, func()) {
+func CreateBackendServerForTestFromConfig(t *testing.T, config *goconf.ConfigFile) (*goconf.ConfigFile, *BackendServer, NatsClient, *Hub, *mux.Router, *httptest.Server) {
 	r := mux.NewRouter()
 	registerBackendHandler(t, r)
 
@@ -103,7 +103,7 @@ func CreateBackendServerForTestFromConfig(t *testing.T, config *goconf.ConfigFil
 
 	go hub.Run()
 
-	shutdown := func() {
+	t.Cleanup(func() {
 		ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 		defer cancel()
 
@@ -111,9 +111,9 @@ func CreateBackendServerForTestFromConfig(t *testing.T, config *goconf.ConfigFil
 		(nats).(*LoopbackNatsClient).waitForSubscriptionsEmpty(ctx, t)
 		nats.Close()
 		server.Close()
-	}
+	})
 
-	return config, b, nats, hub, r, server, shutdown
+	return config, b, nats, hub, r, server
 }
 
 func performBackendRequest(url string, body []byte) (*http.Response, error) {
@@ -161,8 +161,7 @@ func expectRoomlistEvent(n NatsClient, ch chan *nats.Msg, subject string, msgTyp
 }
 
 func TestBackendServer_NoAuth(t *testing.T) {
-	_, _, _, _, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, _, _, server := CreateBackendServerForTest(t)
 
 	roomId := "the-room-id"
 	data := []byte{'{', '}'}
@@ -188,8 +187,7 @@ func TestBackendServer_NoAuth(t *testing.T) {
 }
 
 func TestBackendServer_InvalidAuth(t *testing.T) {
-	_, _, _, _, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, _, _, server := CreateBackendServerForTest(t)
 
 	roomId := "the-room-id"
 	data := []byte{'{', '}'}
@@ -217,8 +215,7 @@ func TestBackendServer_InvalidAuth(t *testing.T) {
 }
 
 func TestBackendServer_OldCompatAuth(t *testing.T) {
-	_, _, _, _, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, _, _, server := CreateBackendServerForTest(t)
 
 	roomId := "the-room-id"
 	userid := "the-user-id"
@@ -267,8 +264,7 @@ func TestBackendServer_OldCompatAuth(t *testing.T) {
 }
 
 func TestBackendServer_InvalidBody(t *testing.T) {
-	_, _, _, _, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, _, _, server := CreateBackendServerForTest(t)
 
 	roomId := "the-room-id"
 	data := []byte{1, 2, 3, 4} // Invalid JSON
@@ -287,8 +283,7 @@ func TestBackendServer_InvalidBody(t *testing.T) {
 }
 
 func TestBackendServer_UnsupportedRequest(t *testing.T) {
-	_, _, _, _, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, _, _, server := CreateBackendServerForTest(t)
 
 	msg := &BackendServerRoomRequest{
 		Type: "lala",
@@ -314,8 +309,7 @@ func TestBackendServer_UnsupportedRequest(t *testing.T) {
 }
 
 func TestBackendServer_RoomInvite(t *testing.T) {
-	_, _, n, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, n, hub, _, server := CreateBackendServerForTest(t)
 
 	u, err := url.Parse(server.URL)
 	if err != nil {
@@ -382,8 +376,7 @@ func TestBackendServer_RoomInvite(t *testing.T) {
 }
 
 func TestBackendServer_RoomDisinvite(t *testing.T) {
-	_, _, n, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, n, hub, _, server := CreateBackendServerForTest(t)
 
 	u, err := url.Parse(server.URL)
 	if err != nil {
@@ -491,8 +484,7 @@ func TestBackendServer_RoomDisinvite(t *testing.T) {
 }
 
 func TestBackendServer_RoomDisinviteDifferentRooms(t *testing.T) {
-	_, _, _, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, hub, _, server := CreateBackendServerForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -618,8 +610,7 @@ func TestBackendServer_RoomDisinviteDifferentRooms(t *testing.T) {
 }
 
 func TestBackendServer_RoomUpdate(t *testing.T) {
-	_, _, n, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, n, hub, _, server := CreateBackendServerForTest(t)
 
 	u, err := url.Parse(server.URL)
 	if err != nil {
@@ -704,8 +695,7 @@ func TestBackendServer_RoomUpdate(t *testing.T) {
 }
 
 func TestBackendServer_RoomDelete(t *testing.T) {
-	_, _, n, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, n, hub, _, server := CreateBackendServerForTest(t)
 
 	u, err := url.Parse(server.URL)
 	if err != nil {
@@ -786,8 +776,7 @@ func TestBackendServer_RoomDelete(t *testing.T) {
 }
 
 func TestBackendServer_ParticipantsUpdatePermissions(t *testing.T) {
-	_, _, _, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, hub, _, server := CreateBackendServerForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -901,8 +890,7 @@ func TestBackendServer_ParticipantsUpdatePermissions(t *testing.T) {
 }
 
 func TestBackendServer_ParticipantsUpdateEmptyPermissions(t *testing.T) {
-	_, _, _, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, hub, _, server := CreateBackendServerForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -987,8 +975,7 @@ func TestBackendServer_ParticipantsUpdateEmptyPermissions(t *testing.T) {
 }
 
 func TestBackendServer_ParticipantsUpdateTimeout(t *testing.T) {
-	_, _, _, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, hub, _, server := CreateBackendServerForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -1201,8 +1188,7 @@ func TestBackendServer_ParticipantsUpdateTimeout(t *testing.T) {
 }
 
 func TestBackendServer_RoomMessage(t *testing.T) {
-	_, _, _, hub, _, server, shutdown := CreateBackendServerForTest(t)
-	defer shutdown()
+	_, _, _, hub, _, server := CreateBackendServerForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1267,8 +1253,7 @@ func TestBackendServer_RoomMessage(t *testing.T) {
 }
 
 func TestBackendServer_TurnCredentials(t *testing.T) {
-	_, _, _, _, _, server, shutdown := CreateBackendServerForTestWithTurn(t)
-	defer shutdown()
+	_, _, _, _, _, server := CreateBackendServerForTestWithTurn(t)
 
 	q := make(url.Values)
 	q.Set("service", "turn")
diff --git a/clientsession_test.go b/clientsession_test.go
index f0604ef..427761f 100644
--- a/clientsession_test.go
+++ b/clientsession_test.go
@@ -126,8 +126,7 @@ func Test_permissionsEqual(t *testing.T) {
 }
 
 func TestBandwidth_Client(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	mcu, err := NewTestMCU()
 	if err != nil {
@@ -199,8 +198,7 @@ func TestBandwidth_Client(t *testing.T) {
 }
 
 func TestBandwidth_Backend(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubWithMultipleBackendsForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubWithMultipleBackendsForTest(t)
 
 	u, err := url.Parse(server.URL + "/one")
 	if err != nil {
diff --git a/hub_test.go b/hub_test.go
index 0326236..9305275 100644
--- a/hub_test.go
+++ b/hub_test.go
@@ -97,7 +97,7 @@ func getTestConfigWithMultipleBackends(server *httptest.Server) (*goconf.ConfigF
 	return config, nil
 }
 
-func CreateHubForTestWithConfig(t *testing.T, getConfigFunc func(*httptest.Server) (*goconf.ConfigFile, error)) (*Hub, NatsClient, *mux.Router, *httptest.Server, func()) {
+func CreateHubForTestWithConfig(t *testing.T, getConfigFunc func(*httptest.Server) (*goconf.ConfigFile, error)) (*Hub, NatsClient, *mux.Router, *httptest.Server) {
 	r := mux.NewRouter()
 	registerBackendHandler(t, r)
 
@@ -124,7 +124,7 @@ func CreateHubForTestWithConfig(t *testing.T, getConfigFunc func(*httptest.Serve
 
 	go h.Run()
 
-	shutdown := func() {
+	t.Cleanup(func() {
 		ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 		defer cancel()
 
@@ -132,20 +132,20 @@ func CreateHubForTestWithConfig(t *testing.T, getConfigFunc func(*httptest.Serve
 		(nats).(*LoopbackNatsClient).waitForSubscriptionsEmpty(ctx, t)
 		nats.Close()
 		server.Close()
-	}
+	})
 
-	return h, nats, r, server, shutdown
+	return h, nats, r, server
 }
 
-func CreateHubForTest(t *testing.T) (*Hub, NatsClient, *mux.Router, *httptest.Server, func()) {
+func CreateHubForTest(t *testing.T) (*Hub, NatsClient, *mux.Router, *httptest.Server) {
 	return CreateHubForTestWithConfig(t, getTestConfig)
 }
 
-func CreateHubWithMultipleBackendsForTest(t *testing.T) (*Hub, NatsClient, *mux.Router, *httptest.Server, func()) {
-	h, nats, r, server, shutdown := CreateHubForTestWithConfig(t, getTestConfigWithMultipleBackends)
+func CreateHubWithMultipleBackendsForTest(t *testing.T) (*Hub, NatsClient, *mux.Router, *httptest.Server) {
+	h, nats, r, server := CreateHubForTestWithConfig(t, getTestConfigWithMultipleBackends)
 	registerBackendHandlerUrl(t, r, "/one")
 	registerBackendHandlerUrl(t, r, "/two")
-	return h, nats, r, server, shutdown
+	return h, nats, r, server
 }
 
 func WaitForHub(ctx context.Context, t *testing.T, h *Hub) {
@@ -426,8 +426,7 @@ func performHousekeeping(hub *Hub, now time.Time) *sync.WaitGroup {
 }
 
 func TestExpectClientHello(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	// The server will send an error and close the connection if no "Hello"
 	// is sent.
@@ -461,8 +460,7 @@ func TestExpectClientHello(t *testing.T) {
 }
 
 func TestClientHello(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -487,8 +485,7 @@ func TestClientHello(t *testing.T) {
 }
 
 func TestClientHelloWithSpaces(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -514,7 +511,7 @@ func TestClientHelloWithSpaces(t *testing.T) {
 }
 
 func TestClientHelloAllowAll(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTestWithConfig(t, func(server *httptest.Server) (*goconf.ConfigFile, error) {
+	hub, _, _, server := CreateHubForTestWithConfig(t, func(server *httptest.Server) (*goconf.ConfigFile, error) {
 		config, err := getTestConfig(server)
 		if err != nil {
 			return nil, err
@@ -524,7 +521,6 @@ func TestClientHelloAllowAll(t *testing.T) {
 		config.AddOption("backend", "allowall", "true")
 		return config, nil
 	})
-	defer shutdown()
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -549,7 +545,7 @@ func TestClientHelloAllowAll(t *testing.T) {
 }
 
 func TestClientHelloSessionLimit(t *testing.T) {
-	hub, _, router, server, shutdown := CreateHubForTestWithConfig(t, func(server *httptest.Server) (*goconf.ConfigFile, error) {
+	hub, _, router, server := CreateHubForTestWithConfig(t, func(server *httptest.Server) (*goconf.ConfigFile, error) {
 		config, err := getTestConfig(server)
 		if err != nil {
 			return nil, err
@@ -567,7 +563,6 @@ func TestClientHelloSessionLimit(t *testing.T) {
 		config.AddOption("backend2", "secret", string(testBackendSecret))
 		return config, nil
 	})
-	defer shutdown()
 
 	registerBackendHandlerUrl(t, router, "/one")
 	registerBackendHandlerUrl(t, router, "/two")
@@ -663,8 +658,7 @@ func TestClientHelloSessionLimit(t *testing.T) {
 }
 
 func TestSessionIdsUnordered(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	publicSessionIds := make([]string, 0)
 	for i := 0; i < 20; i++ {
@@ -737,8 +731,7 @@ func TestSessionIdsUnordered(t *testing.T) {
 }
 
 func TestClientHelloResume(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -793,8 +786,7 @@ func TestClientHelloResume(t *testing.T) {
 }
 
 func TestClientHelloResumeExpired(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -849,8 +841,7 @@ func TestClientHelloResumeExpired(t *testing.T) {
 }
 
 func TestClientHelloResumeTakeover(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -918,8 +909,7 @@ func TestClientHelloResumeTakeover(t *testing.T) {
 }
 
 func TestClientHelloResumeOtherHub(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1014,8 +1004,7 @@ func TestClientHelloResumeOtherHub(t *testing.T) {
 
 func TestClientHelloResumePublicId(t *testing.T) {
 	// Test that a client can't resume a "public" session of another user.
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -1088,8 +1077,7 @@ func TestClientHelloResumePublicId(t *testing.T) {
 }
 
 func TestClientHelloByeResume(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1154,8 +1142,7 @@ func TestClientHelloByeResume(t *testing.T) {
 }
 
 func TestClientHelloResumeAndJoin(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1218,8 +1205,7 @@ func TestClientHelloResumeAndJoin(t *testing.T) {
 }
 
 func TestClientHelloClient(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1247,8 +1233,7 @@ func TestClientHelloClient(t *testing.T) {
 }
 
 func TestClientHelloClient_V3Api(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1281,8 +1266,7 @@ func TestClientHelloClient_V3Api(t *testing.T) {
 }
 
 func TestClientHelloInternal(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1310,8 +1294,7 @@ func TestClientHelloInternal(t *testing.T) {
 }
 
 func TestClientMessageToSessionId(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -1368,8 +1351,7 @@ func TestClientMessageToSessionId(t *testing.T) {
 }
 
 func TestClientMessageToUserId(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -1429,8 +1411,7 @@ func TestClientMessageToUserId(t *testing.T) {
 }
 
 func TestClientMessageToUserIdMultipleSessions(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -1552,8 +1533,7 @@ func WaitForUsersJoined(ctx context.Context, t *testing.T, client1 *TestClient,
 }
 
 func TestClientMessageToRoom(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()
@@ -1627,8 +1607,7 @@ func TestClientMessageToRoom(t *testing.T) {
 }
 
 func TestJoinRoom(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1667,8 +1646,7 @@ func TestJoinRoom(t *testing.T) {
 }
 
 func TestExpectAnonymousJoinRoom(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1720,8 +1698,7 @@ func TestExpectAnonymousJoinRoom(t *testing.T) {
 }
 
 func TestJoinRoomChange(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1773,8 +1750,7 @@ func TestJoinRoomChange(t *testing.T) {
 }
 
 func TestJoinMultiple(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -1852,8 +1828,7 @@ func TestJoinMultiple(t *testing.T) {
 }
 
 func TestJoinRoomSwitchClient(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client := NewTestClient(t, server, hub)
 	defer client.CloseWithBye()
@@ -1978,8 +1953,7 @@ func TestGetRealUserIP(t *testing.T) {
 }
 
 func TestClientMessageToSessionIdWhileDisconnected(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -2070,8 +2044,7 @@ func TestClientMessageToSessionIdWhileDisconnected(t *testing.T) {
 }
 
 func TestRoomParticipantsListUpdateWhileDisconnected(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -2203,8 +2176,7 @@ func TestRoomParticipantsListUpdateWhileDisconnected(t *testing.T) {
 }
 
 func TestClientTakeoverRoomSession(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	client1 := NewTestClient(t, server, hub)
 	defer client1.CloseWithBye()
@@ -2329,8 +2301,7 @@ func TestClientTakeoverRoomSession(t *testing.T) {
 }
 
 func TestClientSendOfferPermissions(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	mcu, err := NewTestMCU()
 	if err != nil {
@@ -2457,8 +2428,7 @@ func TestClientSendOfferPermissions(t *testing.T) {
 }
 
 func TestClientSendOfferPermissionsAudioOnly(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	mcu, err := NewTestMCU()
 	if err != nil {
@@ -2549,8 +2519,7 @@ func TestClientSendOfferPermissionsAudioOnly(t *testing.T) {
 }
 
 func TestClientSendOfferPermissionsAudioVideo(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	mcu, err := NewTestMCU()
 	if err != nil {
@@ -2678,8 +2647,7 @@ loop:
 }
 
 func TestClientSendOfferPermissionsAudioVideoMedia(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	mcu, err := NewTestMCU()
 	if err != nil {
@@ -2811,8 +2779,7 @@ loop:
 }
 
 func TestClientRequestOfferNotInRoom(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	mcu, err := NewTestMCU()
 	if err != nil {
@@ -2997,8 +2964,7 @@ func TestClientRequestOfferNotInRoom(t *testing.T) {
 
 func TestNoSendBetweenSessionsOnDifferentBackends(t *testing.T) {
 	// Clients can't send messages to sessions connected from other backends.
-	hub, _, _, server, shutdown := CreateHubWithMultipleBackendsForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubWithMultipleBackendsForTest(t)
 
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()
@@ -3068,8 +3034,7 @@ func TestNoSendBetweenSessionsOnDifferentBackends(t *testing.T) {
 }
 
 func TestNoSameRoomOnDifferentBackends(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubWithMultipleBackendsForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubWithMultipleBackendsForTest(t)
 
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()
diff --git a/natsclient_loopback_test.go b/natsclient_loopback_test.go
index 99aad5b..3cb35b1 100644
--- a/natsclient_loopback_test.go
+++ b/natsclient_loopback_test.go
@@ -48,20 +48,20 @@ func (c *LoopbackNatsClient) waitForSubscriptionsEmpty(ctx context.Context, t *t
 	}
 }
 
-func CreateLoopbackNatsClientForTest(t *testing.T) (NatsClient, func()) {
+func CreateLoopbackNatsClientForTest(t *testing.T) NatsClient {
 	result, err := NewLoopbackNatsClient()
 	if err != nil {
 		t.Fatal(err)
 	}
-	return result, func() {
+	t.Cleanup(func() {
 		result.Close()
-	}
+	})
+	return result
 }
 
 func TestLoopbackNatsClient_Subscribe(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLoopbackNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLoopbackNatsClientForTest(t)
 
 		testNatsClient_Subscribe(t, client)
 	})
@@ -69,8 +69,7 @@ func TestLoopbackNatsClient_Subscribe(t *testing.T) {
 
 func TestLoopbackClient_PublishAfterClose(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLoopbackNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLoopbackNatsClientForTest(t)
 
 		testNatsClient_PublishAfterClose(t, client)
 	})
@@ -78,8 +77,7 @@ func TestLoopbackClient_PublishAfterClose(t *testing.T) {
 
 func TestLoopbackClient_SubscribeAfterClose(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLoopbackNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLoopbackNatsClientForTest(t)
 
 		testNatsClient_SubscribeAfterClose(t, client)
 	})
@@ -87,8 +85,7 @@ func TestLoopbackClient_SubscribeAfterClose(t *testing.T) {
 
 func TestLoopbackClient_BadSubjects(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLoopbackNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLoopbackNatsClientForTest(t)
 
 		testNatsClient_BadSubjects(t, client)
 	})
diff --git a/natsclient_test.go b/natsclient_test.go
index ee45b61..62defca 100644
--- a/natsclient_test.go
+++ b/natsclient_test.go
@@ -31,28 +31,28 @@ import (
 	natsserver "github.com/nats-io/nats-server/v2/test"
 )
 
-func startLocalNatsServer() (string, func()) {
+func startLocalNatsServer(t *testing.T) string {
 	opts := natsserver.DefaultTestOptions
 	opts.Port = -1
 	opts.Cluster.Name = "testing"
 	srv := natsserver.RunServer(&opts)
-	shutdown := func() {
+	t.Cleanup(func() {
 		srv.Shutdown()
 		srv.WaitForShutdown()
-	}
-	return srv.ClientURL(), shutdown
+	})
+	return srv.ClientURL()
 }
 
-func CreateLocalNatsClientForTest(t *testing.T) (NatsClient, func()) {
-	url, shutdown := startLocalNatsServer()
+func CreateLocalNatsClientForTest(t *testing.T) NatsClient {
+	url := startLocalNatsServer(t)
 	result, err := NewNatsClient(url)
 	if err != nil {
 		t.Fatal(err)
 	}
-	return result, func() {
+	t.Cleanup(func() {
 		result.Close()
-		shutdown()
-	}
+	})
+	return result
 }
 
 func testNatsClient_Subscribe(t *testing.T, client NatsClient) {
@@ -106,8 +106,7 @@ func testNatsClient_Subscribe(t *testing.T, client NatsClient) {
 
 func TestNatsClient_Subscribe(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLocalNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLocalNatsClientForTest(t)
 
 		testNatsClient_Subscribe(t, client)
 	})
@@ -123,8 +122,7 @@ func testNatsClient_PublishAfterClose(t *testing.T, client NatsClient) {
 
 func TestNatsClient_PublishAfterClose(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLocalNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLocalNatsClientForTest(t)
 
 		testNatsClient_PublishAfterClose(t, client)
 	})
@@ -141,8 +139,7 @@ func testNatsClient_SubscribeAfterClose(t *testing.T, client NatsClient) {
 
 func TestNatsClient_SubscribeAfterClose(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLocalNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLocalNatsClientForTest(t)
 
 		testNatsClient_SubscribeAfterClose(t, client)
 	})
@@ -164,8 +161,7 @@ func testNatsClient_BadSubjects(t *testing.T, client NatsClient) {
 
 func TestNatsClient_BadSubjects(t *testing.T) {
 	ensureNoGoroutinesLeak(t, func() {
-		client, shutdown := CreateLocalNatsClientForTest(t)
-		defer shutdown()
+		client := CreateLocalNatsClientForTest(t)
 
 		testNatsClient_BadSubjects(t, client)
 	})
diff --git a/proxy/proxy_server_test.go b/proxy/proxy_server_test.go
index 62e6072..70a128f 100644
--- a/proxy/proxy_server_test.go
+++ b/proxy/proxy_server_test.go
@@ -41,14 +41,14 @@ const (
 	TokenIdForTest     = "foo"
 )
 
-func newProxyServerForTest(t *testing.T) (*ProxyServer, *rsa.PrivateKey, func()) {
+func newProxyServerForTest(t *testing.T) (*ProxyServer, *rsa.PrivateKey) {
 	tempdir := t.TempDir()
 	var server *ProxyServer
-	shutdown := func() {
+	t.Cleanup(func() {
 		if server != nil {
 			server.Stop()
 		}
-	}
+	})
 
 	r := mux.NewRouter()
 	key, err := rsa.GenerateKey(rand.Reader, KeypairSizeForTest)
@@ -89,12 +89,11 @@ func newProxyServerForTest(t *testing.T) (*ProxyServer, *rsa.PrivateKey, func())
 	if server, err = NewProxyServer(r, "0.0", config); err != nil {
 		t.Fatalf("could not create server: %s", err)
 	}
-	return server, key, shutdown
+	return server, key
 }
 
 func TestTokenInFuture(t *testing.T) {
-	server, key, shutdown := newProxyServerForTest(t)
-	defer shutdown()
+	server, key := newProxyServerForTest(t)
 
 	claims := &signaling.TokenClaims{
 		StandardClaims: jwt.StandardClaims{
diff --git a/room_test.go b/room_test.go
index 49f3dcb..ab7caf4 100644
--- a/room_test.go
+++ b/room_test.go
@@ -73,8 +73,7 @@ func TestRoom_InCall(t *testing.T) {
 }
 
 func TestRoom_Update(t *testing.T) {
-	hub, _, router, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, router, server := CreateHubForTest(t)
 
 	config, err := getTestConfig(server)
 	if err != nil {
@@ -211,8 +210,7 @@ loop:
 }
 
 func TestRoom_Delete(t *testing.T) {
-	hub, _, router, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, router, server := CreateHubForTest(t)
 
 	config, err := getTestConfig(server)
 	if err != nil {
@@ -354,8 +352,7 @@ loop:
 }
 
 func TestRoom_RoomSessionData(t *testing.T) {
-	hub, _, router, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, router, server := CreateHubForTest(t)
 
 	config, err := getTestConfig(server)
 	if err != nil {
@@ -424,8 +421,7 @@ func TestRoom_RoomSessionData(t *testing.T) {
 }
 
 func TestRoom_InCallAll(t *testing.T) {
-	hub, _, router, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, router, server := CreateHubForTest(t)
 
 	config, err := getTestConfig(server)
 	if err != nil {
diff --git a/transient_data_test.go b/transient_data_test.go
index 43c99c2..fb8e89f 100644
--- a/transient_data_test.go
+++ b/transient_data_test.go
@@ -68,8 +68,7 @@ func Test_TransientData(t *testing.T) {
 }
 
 func Test_TransientMessages(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()
diff --git a/virtualsession_test.go b/virtualsession_test.go
index e6dac71..1c8890e 100644
--- a/virtualsession_test.go
+++ b/virtualsession_test.go
@@ -28,8 +28,7 @@ import (
 )
 
 func TestVirtualSession(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	roomId := "the-room-id"
 	emptyProperties := json.RawMessage("{}")
@@ -311,8 +310,7 @@ func TestVirtualSession(t *testing.T) {
 }
 
 func TestVirtualSessionCleanup(t *testing.T) {
-	hub, _, _, server, shutdown := CreateHubForTest(t)
-	defer shutdown()
+	hub, _, _, server := CreateHubForTest(t)
 
 	roomId := "the-room-id"
 	emptyProperties := json.RawMessage("{}")