1 files changed, 8 insertions, 10 deletions
diff --git a/api/websocket_test.go b/api/websocket_test.go
index a5c512262..161ea4960 100644
--- a/api/websocket_test.go
+++ b/api/websocket_test.go
@@ -116,6 +116,8 @@ import (
 
 func TestWebSocket(t *testing.T) {
 	th := Setup().InitBasic()
+	defer th.TearDown()
+
 	WebSocketClient, err := th.CreateWebSocketClient()
 	if err != nil {
 		t.Fatal(err)
@@ -177,6 +179,8 @@ func TestWebSocket(t *testing.T) {
 
 func TestWebSocketEvent(t *testing.T) {
 	th := Setup().InitBasic()
+	defer th.TearDown()
+
 	WebSocketClient, err := th.CreateWebSocketClient()
 	if err != nil {
 		t.Fatal(err)
@@ -252,6 +256,8 @@ func TestWebSocketEvent(t *testing.T) {
 
 func TestCreateDirectChannelWithSocket(t *testing.T) {
 	th := Setup().InitBasic()
+	defer th.TearDown()
+
 	Client := th.BasicClient
 	user2 := th.BasicUser2
 
@@ -314,7 +320,8 @@ func TestCreateDirectChannelWithSocket(t *testing.T) {
 }
 
 func TestWebsocketOriginSecurity(t *testing.T) {
-	Setup().InitBasic()
+	th := Setup().InitBasic()
+	defer th.TearDown()
 
 	url := "ws://localhost" + *utils.Cfg.ServiceSettings.ListenAddress
 
@@ -372,12 +379,3 @@ func TestWebsocketOriginSecurity(t *testing.T) {
 
 	*utils.Cfg.ServiceSettings.AllowCorsFrom = ""
 }
-
-func TestZZWebSocketTearDown(t *testing.T) {
-	// *IMPORTANT* - Kind of hacky
-	// This should be the last function in any test file
-	// that calls Setup()
-	// Should be in the last file too sorted by name
-	time.Sleep(2 * time.Second)
-	TearDown()
-}