Use CreateConnectedWebSocketClient test util (#30036)

server/channels/api4/apitestlib.go

@@ -78,7 +78,8 @@ func SetMainHelper(mh *testlib.MainHelper) {
 }
 
 func setupTestHelper(dbStore store.Store, searchEngine *searchengine.Broker, enterprise bool, includeCache bool,
-	updateConfig func(*model.Config), options []app.Option) *TestHelper {
+	updateConfig func(*model.Config), options []app.Option,
+) *TestHelper {
 	tempWorkspace, err := os.MkdirTemp("", "apptest")
 	if err != nil {
 		panic(err)
@@ -421,14 +422,16 @@ func closeBody(r *http.Response) {
 	}
 }
 
-var initBasicOnce sync.Once
+var (
-var userCache struct {
+	initBasicOnce sync.Once
-	SystemAdminUser   *model.User
+	userCache     struct {
-	SystemManagerUser *model.User
+		SystemAdminUser   *model.User
-	TeamAdminUser     *model.User
+		SystemManagerUser *model.User
-	BasicUser         *model.User
+		TeamAdminUser     *model.User
-	BasicUser2        *model.User
+		BasicUser         *model.User
-}
+		BasicUser2        *model.User
+	}
+)
 
 func (th *TestHelper) InitLogin() *TestHelper {
 	th.waitForConnectivity()
@@ -555,9 +558,9 @@ func (th *TestHelper) CreateLocalClient(socketPath string) *model.Client4 {
 	}
 }
 
-func (th *TestHelper) CreateConnectedWebSocketClient(t *testing.T) *model.WebSocketClient {
+func (th *TestHelper) createConnectedWebSocketClient(t *testing.T, client *model.Client4) *model.WebSocketClient {
 	t.Helper()
-	wsClient, err := th.CreateWebSocketClient()
+	wsClient, err := th.CreateWebSocketClientWithClient(client)
 	require.NoError(t, err)
 	require.NotNil(t, wsClient, "webSocketClient should not be nil")
 	wsClient.Listen()
@@ -574,6 +577,14 @@ func (th *TestHelper) CreateConnectedWebSocketClient(t *testing.T) *model.WebSoc
 	return wsClient
 }
 
+func (th *TestHelper) CreateConnectedWebSocketClient(t *testing.T) *model.WebSocketClient {
+	return th.createConnectedWebSocketClient(t, th.Client)
+}
+
+func (th *TestHelper) CreateConnectedWebSocketClientWithClient(t *testing.T, client *model.Client4) *model.WebSocketClient {
+	return th.createConnectedWebSocketClient(t, client)
+}
+
 func (th *TestHelper) CreateWebSocketClient() (*model.WebSocketClient, error) {
 	return model.NewWebSocketClient4(fmt.Sprintf("ws://localhost:%v", th.App.Srv().ListenAddr.Port), th.Client.AuthToken)
 }
@@ -582,10 +593,6 @@ func (th *TestHelper) CreateReliableWebSocketClient(connID string, seqNo int) (*
 	return model.NewReliableWebSocketClientWithDialer(websocket.DefaultDialer, fmt.Sprintf("ws://localhost:%v", th.App.Srv().ListenAddr.Port), th.Client.AuthToken, connID, seqNo, true)
 }
 
-func (th *TestHelper) CreateWebSocketSystemAdminClient() (*model.WebSocketClient, error) {
-	return model.NewWebSocketClient4(fmt.Sprintf("ws://localhost:%v", th.App.Srv().ListenAddr.Port), th.SystemAdminClient.AuthToken)
-}
-
 func (th *TestHelper) CreateWebSocketClientWithClient(client *model.Client4) (*model.WebSocketClient, error) {
 	return model.NewWebSocketClient4(fmt.Sprintf("ws://localhost:%v", th.App.Srv().ListenAddr.Port), client.AuthToken)
 }
@@ -877,7 +884,6 @@ func (th *TestHelper) CreateMessagePostNoClient(channel *model.Channel, message
 		Message:   message,
 		CreateAt:  createAtTime,
 	})
-
 	if err != nil {
 		panic(err)
 	}

server/channels/api4/channel_category_test.go

@@ -106,11 +106,7 @@ func TestCreateCategoryForTeamForUser(t *testing.T) {
 	})
 
 	t.Run("should publish expected WS payload", func(t *testing.T) {
-		t.Skip("MM-42652")
+		userWSClient := th.CreateConnectedWebSocketClient(t)
-		userWSClient, err := th.CreateWebSocketClient()
-		require.NoError(t, err)
-		defer userWSClient.Close()
-		userWSClient.Listen()
 
 		category := &model.SidebarCategoryWithChannels{
 			SidebarCategory: model.SidebarCategory{

server/channels/api4/channel_test.go

@@ -11,7 +11,6 @@ import (
 	"net/http"
 	"sort"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 
@@ -175,7 +174,7 @@ func TestUpdateChannel(t *testing.T) {
 	channel, _, _ = client.CreateChannel(context.Background(), channel)
 	private, _, _ = client.CreateChannel(context.Background(), private)
 
-	//Update a open channel
+	// Update a open channel
 	channel.DisplayName = "My new display name"
 	channel.Header = "My fancy header"
 	channel.Purpose = "Mattermost ftw!"
@@ -195,7 +194,7 @@ func TestUpdateChannel(t *testing.T) {
 
 	require.Equal(t, *channel.GroupConstrained, *rchannel.GroupConstrained, "GroupConstrained flags do not match")
 
-	//Update a private channel
+	// Update a private channel
 	private.DisplayName = "My new display name for private channel"
 	private.Header = "My fancy private header"
 	private.Purpose = "Mattermost ftw! in private mode"
@@ -207,7 +206,7 @@ func TestUpdateChannel(t *testing.T) {
 	require.Equal(t, private.Header, newPrivateChannel.Header, "Update failed for Header in private channel")
 	require.Equal(t, private.Purpose, newPrivateChannel.Purpose, "Update failed for Purpose in private channel")
 
-	//Test updating default channel's name and returns error
+	// Test updating default channel's name and returns error
 	defaultChannel, _ := th.App.GetChannelByName(th.Context, model.DefaultChannelName, team.Id, false)
 	defaultChannel.Name = "testing"
 	_, resp, err = client.UpdateChannel(context.Background(), defaultChannel)
@@ -226,19 +225,19 @@ func TestUpdateChannel(t *testing.T) {
 	_, _, err = client.UpdateChannel(context.Background(), private)
 	require.NoError(t, err)
 
-	//Non existing channel
+	// Non existing channel
 	channel1 := &model.Channel{DisplayName: "Test API Name for apiv4", Name: GenerateTestChannelName(), Type: model.ChannelTypeOpen, TeamId: team.Id}
 	_, resp, err = client.UpdateChannel(context.Background(), channel1)
 	require.Error(t, err)
 	CheckNotFoundStatus(t, resp)
 
-	//Try to update with not logged user
+	// Try to update with not logged user
 	client.Logout(context.Background())
 	_, resp, err = client.UpdateChannel(context.Background(), channel)
 	require.Error(t, err)
 	CheckUnauthorizedStatus(t, resp)
 
-	//Try to update using another user
+	// Try to update using another user
 	user := th.CreateUser()
 	client.Login(context.Background(), user.Email, user.Password)
 
@@ -376,7 +375,7 @@ func TestPatchChannel(t *testing.T) {
 
 	require.Equal(t, oldName, channel.Name, "should not have updated")
 
-	//Test updating default channel's name and returns error
+	// Test updating default channel's name and returns error
 	defaultChannel, _ := th.App.GetChannelByName(th.Context, model.DefaultChannelName, team.Id, false)
 	defaultChannelPatch := &model.ChannelPatch{
 		Name: new(string),
@@ -532,7 +531,8 @@ func TestChannelUnicodeNames(t *testing.T) {
 			Name:        "\u206cenglish\u206dchannel",
 			DisplayName: "The \u206cEnglish\u206d Channel",
 			Type:        model.ChannelTypeOpen,
-			TeamId:      team.Id}
+			TeamId:      team.Id,
+		}
 
 		rchannel, resp, err := client.CreateChannel(context.Background(), channel)
 		require.NoError(t, err)
@@ -3137,9 +3137,7 @@ func TestUpdateChannelMemberSchemeRoles(t *testing.T) {
 	require.Nil(t, appError)
 
 	SystemAdminClient := th.SystemAdminClient
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	WebSocketClient.Listen()
-	require.NoError(t, err)
 
 	th.LoginBasic()
 
@@ -3148,7 +3146,7 @@ func TestUpdateChannelMemberSchemeRoles(t *testing.T) {
 		SchemeUser:  false,
 		SchemeGuest: false,
 	}
-	_, err = SystemAdminClient.UpdateChannelMemberSchemeRoles(context.Background(), th.BasicChannel.Id, th.BasicUser.Id, s1)
+	_, err := SystemAdminClient.UpdateChannelMemberSchemeRoles(context.Background(), th.BasicChannel.Id, th.BasicUser.Id, s1)
 	require.NoError(t, err)
 
 	waiting := true
@@ -3185,7 +3183,7 @@ func TestUpdateChannelMemberSchemeRoles(t *testing.T) {
 	assert.Equal(t, true, tm2.SchemeUser)
 	assert.Equal(t, false, tm2.SchemeAdmin)
 
-	//cannot set Guest to User for single channel
+	// cannot set Guest to User for single channel
 	resp, err := SystemAdminClient.UpdateChannelMemberSchemeRoles(context.Background(), th.BasicChannel.Id, guest.Id, s2)
 	require.Error(t, err)
 	CheckBadRequestStatus(t, resp)
@@ -3568,10 +3566,7 @@ func TestAddChannelMemberFromThread(t *testing.T) {
 	_, _, err := th.SystemAdminClient.AddTeamMember(context.Background(), team.Id, user3.Id)
 	require.NoError(t, err)
 
-	wsClient, err2 := th.CreateWebSocketClient()
+	wsClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err2)
-	defer wsClient.Close()
-	wsClient.Listen()
 
 	publicChannel := th.CreatePublicChannel()
 
@@ -3767,16 +3762,7 @@ func TestRemoveChannelMember(t *testing.T) {
 		_, err = th.SystemAdminClient.UpdateChannelNotifyProps(context.Background(), th.BasicChannel2.Id, th.SystemAdminUser.Id, props)
 		require.NoError(t, err)
 
-		wsClient, err2 := th.CreateWebSocketSystemAdminClient()
+		wsClient := th.CreateConnectedWebSocketClientWithClient(t, th.SystemAdminClient)
-		require.NoError(t, err2)
-		wsClient.Listen()
-		var closeWsClient sync.Once
-		defer closeWsClient.Do(func() {
-			wsClient.Close()
-		})
-
-		wsr := <-wsClient.EventChannel
-		require.Equal(t, model.WebsocketEventHello, wsr.EventType())
 
 		// requirePost listens for websocket events and tries to find the post matching
 		// the expected post's channel and message.
@@ -3803,7 +3789,7 @@ func TestRemoveChannelMember(t *testing.T) {
 		}
 
 		th.App.AddUserToChannel(th.Context, th.BasicUser2, th.BasicChannel, false)
-		_, err2 = client.RemoveUserFromChannel(context.Background(), th.BasicChannel.Id, th.BasicUser2.Id)
+		_, err2 := client.RemoveUserFromChannel(context.Background(), th.BasicChannel.Id, th.BasicUser2.Id)
 		require.NoError(t, err2)
 
 		requirePost(&model.Post{
@@ -3824,10 +3810,6 @@ func TestRemoveChannelMember(t *testing.T) {
 			Message:   fmt.Sprintf("@%s removed from the channel.", th.BasicUser.Username),
 			ChannelId: th.BasicChannel.Id,
 		})
-
-		closeWsClient.Do(func() {
-			wsClient.Close()
-		})
 	})
 
 	// Leave deleted channel
@@ -4365,7 +4347,7 @@ func TestGetChannelMembersTimezones(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, timezone, 2, "should return 2 timezones")
 
-	//both users have same timezone
+	// both users have same timezone
 	user2.Timezone["automaticTimezone"] = "XOXO/BLABLA"
 	_, _, err = th.SystemAdminClient.UpdateUser(context.Background(), user2)
 	require.NoError(t, err)
@@ -4374,7 +4356,7 @@ func TestGetChannelMembersTimezones(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, timezone, 1, "should return 1 timezone")
 
-	//no timezone set should return empty
+	// no timezone set should return empty
 	user2.Timezone["automaticTimezone"] = ""
 	_, _, err = th.SystemAdminClient.UpdateUser(context.Background(), user2)
 	require.NoError(t, err)

server/channels/api4/plugin_test.go

@@ -336,10 +336,8 @@ func TestNotifyClusterPluginEvent(t *testing.T) {
 	require.Equal(t, "testplugin", manifest.Id)
 
 	// Successful remove
-	webSocketClient, err := th.CreateWebSocketSystemAdminClient()
+	webSocketClient := th.CreateConnectedWebSocketClientWithClient(t, th.SystemAdminClient)
-	require.NoError(t, err)
+
-	webSocketClient.Listen()
-	defer webSocketClient.Close()
 	done := make(chan bool)
 	go func() {
 		for {

server/channels/api4/post_test.go

@@ -668,7 +668,7 @@ func testCreatePostWithOutgoingHook(
 			}
 		}
 
-		respPostType := "" //if is empty or post will do a normal post.
+		respPostType := "" // if is empty or post will do a normal post.
 		if commentPostType {
 			respPostType = model.OutgoingHookResponseTypeComment
 		}
@@ -1046,8 +1046,10 @@ func TestMoveThread(t *testing.T) {
 		// Set permitted role as channel admin
 		enabled := true
 		th.App.UpdateConfig(func(cfg *model.Config) {
-			cfg.WranglerSettings = model.WranglerSettings{MoveThreadToAnotherTeamEnable: &enabled,
+			cfg.WranglerSettings = model.WranglerSettings{
-				PermittedWranglerRoles: []string{model.PermissionsChannelAdmin}}
+				MoveThreadToAnotherTeamEnable: &enabled,
+				PermittedWranglerRoles:        []string{model.PermissionsChannelAdmin},
+			}
 		})
 		defer th.App.UpdateConfig(func(cfg *model.Config) {
 			cfg.WranglerSettings = model.WranglerSettings{}
@@ -1215,9 +1217,7 @@ func TestCreatePostSendOutOfChannelMentions(t *testing.T) {
 	defer th.TearDown()
 	client := th.Client
 
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	WebSocketClient.Listen()
 
 	inChannelUser := th.CreateUser()
 	th.LinkUserToTeam(inChannelUser, th.BasicTeam)
@@ -1286,11 +1286,7 @@ func TestCreatePostCheckOnlineStatus(t *testing.T) {
 	_, _, err = cli.Login(context.Background(), th.BasicUser2.Username, th.BasicUser2.Password)
 	require.NoError(t, err)
 
-	wsClient, err := th.CreateWebSocketClientWithClient(cli)
+	wsClient := th.CreateConnectedWebSocketClientWithClient(t, cli)
-	require.NoError(t, err)
-	defer wsClient.Close()
-
-	wsClient.Listen()
 
 	waitForEvent := func(isSetOnline bool) {
 		timeout := time.After(5 * time.Second)
@@ -3394,11 +3390,7 @@ func TestWebHubMembership(t *testing.T) {
 		_, _, err := cli.Login(context.Background(), obj.user.Username, obj.user.Password)
 		require.NoError(t, err)
 
-		wsClient, err := th.CreateWebSocketClientWithClient(cli)
+		wsClient := th.CreateConnectedWebSocketClientWithClient(t, cli)
-		require.NoError(t, err)
-		defer wsClient.Close()
-
-		wsClient.Listen()
 
 		go func(testName string) {
 			defer wg.Done()
@@ -3509,12 +3501,9 @@ func TestDeletePostEvent(t *testing.T) {
 	th := Setup(t).InitBasic()
 	defer th.TearDown()
 
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	WebSocketClient.Listen()
-	defer WebSocketClient.Close()
 
-	_, err = th.SystemAdminClient.DeletePost(context.Background(), th.BasicPost.Id)
+	_, err := th.SystemAdminClient.DeletePost(context.Background(), th.BasicPost.Id)
 	require.NoError(t, err)
 
 	var received, exit bool
@@ -3554,15 +3543,11 @@ func TestDeletePostMessage(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			wsClient, err := th.CreateWebSocketClientWithClient(tc.client)
+			wsClient := th.CreateConnectedWebSocketClientWithClient(t, tc.client)
-			require.NoError(t, err)
-			defer wsClient.Close()
-
-			wsClient.Listen()
 
 			post := th.CreatePost()
 
-			_, err = th.SystemAdminClient.DeletePost(context.Background(), post.Id)
+			_, err := th.SystemAdminClient.DeletePost(context.Background(), post.Id)
 			require.NoError(t, err)
 
 			timeout := time.After(5 * time.Second)
@@ -4310,12 +4295,9 @@ func TestSetPostUnreadWithoutCollapsedThreads(t *testing.T) {
 	require.Nil(t, appErr)
 
 	t.Run("Mark reply post as unread", func(t *testing.T) {
-		userWSClient, err := th.CreateWebSocketClient()
+		userWSClient := th.CreateConnectedWebSocketClient(t)
-		require.NoError(t, err)
-		defer userWSClient.Close()
-		userWSClient.Listen()
 
-		_, err = th.Client.SetPostUnread(context.Background(), th.BasicUser.Id, replyPost1.Id, false)
+		_, err := th.Client.SetPostUnread(context.Background(), th.BasicUser.Id, replyPost1.Id, false)
 		require.NoError(t, err)
 		channelUnread, appErr := th.App.GetChannelUnread(th.Context, th.BasicChannel.Id, th.BasicUser.Id)
 		require.Nil(t, appErr)
@@ -4375,6 +4357,7 @@ func TestSetPostUnreadWithoutCollapsedThreads(t *testing.T) {
 		require.Equal(t, int64(3), channelUnread.MsgCountRoot)
 	})
 }
+
 func TestGetPostsByIds(t *testing.T) {
 	th := Setup(t).InitBasic()
 	defer th.TearDown()
@@ -4623,10 +4606,7 @@ func TestCreatePostNotificationsWithCRT(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			userWSClient, err := th.CreateWebSocketClient()
+			userWSClient := th.CreateConnectedWebSocketClient(t)
-			require.NoError(t, err)
-			defer userWSClient.Close()
-			userWSClient.Listen()
 
 			patch := &model.UserPatch{}
 			patch.NotifyProps = model.CopyStringMap(th.BasicUser.NotifyProps)
@@ -4635,7 +4615,7 @@ func TestCreatePostNotificationsWithCRT(t *testing.T) {
 			}
 
 			// update user's notify props
-			_, _, err = th.Client.PatchUser(context.Background(), th.BasicUser.Id, patch)
+			_, _, err := th.Client.PatchUser(context.Background(), th.BasicUser.Id, patch)
 			require.NoError(t, err)
 
 			// post a reply on the thread
@@ -4734,10 +4714,7 @@ func TestPostReminder(t *testing.T) {
 	defer th.TearDown()
 
 	client := th.Client
-	userWSClient, err := th.CreateWebSocketClient()
+	userWSClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	defer userWSClient.Close()
-	userWSClient.Listen()
 
 	targetTime := time.Now().UTC().Unix()
 	resp, err := client.SetPostReminder(context.Background(), &model.PostReminder{

server/channels/api4/preference_test.go

@@ -344,13 +344,7 @@ func TestUpdatePreferencesWebsocket(t *testing.T) {
 	th := Setup(t).InitBasic()
 	defer th.TearDown()
 
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-
-	WebSocketClient.Listen()
-	time.Sleep(300 * time.Millisecond)
-	wsResp := <-WebSocketClient.ResponseChannel
-	require.Equal(t, wsResp.Status, model.StatusOk, "expected OK from auth challenge")
 
 	userId := th.BasicUser.Id
 	preferences := model.Preferences{
@@ -366,7 +360,7 @@ func TestUpdatePreferencesWebsocket(t *testing.T) {
 		},
 	}
 
-	_, err = th.Client.UpdatePreferences(context.Background(), userId, preferences)
+	_, err := th.Client.UpdatePreferences(context.Background(), userId, preferences)
 	require.NoError(t, err)
 
 	timeout := time.After(300 * time.Millisecond)
@@ -750,12 +744,7 @@ func TestDeletePreferencesWebsocket(t *testing.T) {
 	_, err := th.Client.UpdatePreferences(context.Background(), userId, preferences)
 	require.NoError(t, err)
 
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-
-	WebSocketClient.Listen()
-	wsResp := <-WebSocketClient.ResponseChannel
-	require.Equal(t, model.StatusOk, wsResp.Status, "should have responded OK to authentication challenge")
 
 	_, err = th.Client.DeletePreferences(context.Background(), userId, preferences)
 	require.NoError(t, err)

server/channels/api4/system_test.go

@@ -590,9 +590,7 @@ func TestGetAnalyticsOld(t *testing.T) {
 	assert.Equal(t, "total_websocket_connections", rows2[5].Name)
 	assert.Equal(t, float64(0), rows2[5].Value)
 
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	time.Sleep(100 * time.Millisecond)
 	rows2, _, err = th.SystemAdminClient.GetAnalyticsOld(context.Background(), "standard", "")
 	require.NoError(t, err)
 	assert.Equal(t, "total_websocket_connections", rows2[5].Name)

server/channels/api4/team_test.go

@@ -915,7 +915,8 @@ func TestTeamUnicodeNames(t *testing.T) {
 			DisplayName: "Some\u206c Team",
 			Description: "A \ufffatest\ufffb channel.",
 			CompanyName: "\ufeffAcme Inc\ufffc",
-			Type:        model.TeamOpen}
+			Type:        model.TeamOpen,
+		}
 		rteam, resp, err := client.CreateTeam(context.Background(), team)
 		require.NoError(t, err)
 		CheckCreatedStatus(t, resp)
@@ -932,7 +933,8 @@ func TestTeamUnicodeNames(t *testing.T) {
 			CompanyName: "Bad Company",
 			Name:        model.NewRandomTeamName(),
 			Email:       "success+" + model.NewId() + "@simulator.amazonses.com",
-			Type:        model.TeamOpen}
+			Type:        model.TeamOpen,
+		}
 		team, _, _ = client.CreateTeam(context.Background(), team)
 
 		team.DisplayName = "\u206eThe Team\u206f"
@@ -953,7 +955,8 @@ func TestTeamUnicodeNames(t *testing.T) {
 			CompanyName: "Some company name",
 			Name:        model.NewRandomTeamName(),
 			Email:       "success+" + model.NewId() + "@simulator.amazonses.com",
-			Type:        model.TeamOpen}
+			Type:        model.TeamOpen,
+		}
 		team, _, _ = client.CreateTeam(context.Background(), team)
 
 		patch := &model.TeamPatch{}
@@ -2814,21 +2817,11 @@ func TestRemoveTeamMemberEvents(t *testing.T) {
 
 	client1 := th.CreateClient()
 	th.LoginBasicWithClient(client1)
-	WebSocketClient, err := th.CreateWebSocketClientWithClient(client1)
+	WebSocketClient := th.CreateConnectedWebSocketClientWithClient(t, client1)
-	require.NoError(t, err)
-	defer WebSocketClient.Close()
-	WebSocketClient.Listen()
-	resp := <-WebSocketClient.ResponseChannel
-	require.Equal(t, resp.Status, model.StatusOk)
 
 	client2 := th.CreateClient()
 	th.LoginBasic2WithClient(client2)
-	WebSocketClient2, err := th.CreateWebSocketClientWithClient(client2)
+	WebSocketClient2 := th.CreateConnectedWebSocketClientWithClient(t, client2)
-	require.NoError(t, err)
-	defer WebSocketClient2.Close()
-	WebSocketClient2.Listen()
-	resp = <-WebSocketClient2.ResponseChannel
-	require.Equal(t, resp.Status, model.StatusOk)
 
 	th.TestForSystemAdminAndLocal(t, func(t *testing.T, client *model.Client4) {
 		// remove second user from basic team
@@ -3039,7 +3032,7 @@ func TestUpdateTeamMemberSchemeRoles(t *testing.T) {
 	assert.Equal(t, true, tm2.SchemeUser)
 	assert.Equal(t, false, tm2.SchemeAdmin)
 
-	//cannot set Guest to User for single team
+	// cannot set Guest to User for single team
 	resp, err := SystemAdminClient.UpdateTeamMemberSchemeRoles(context.Background(), th.BasicTeam.Id, guest.Id, s2)
 	require.Error(t, err)
 	CheckBadRequestStatus(t, resp)
@@ -3381,7 +3374,7 @@ func TestInviteUsersToTeam(t *testing.T) {
 	memberInvite := &model.MemberInvite{Emails: []string{user1, user2}}
 	emailList := memberInvite.Emails
 
-	//Delete all the messages before check the sample email
+	// Delete all the messages before check the sample email
 	err := mail.DeleteMailBox(user1)
 	require.NoError(t, err)
 	err = mail.DeleteMailBox(user2)
@@ -3401,7 +3394,7 @@ func TestInviteUsersToTeam(t *testing.T) {
 	})
 
 	checkEmail := func(t *testing.T, expectedSubject string) {
-		//Check if the email was sent to the right email address
+		// Check if the email was sent to the right email address
 		for _, email := range emailList {
 			var resultsMailbox mail.JSONMessageHeaderInbucket
 			err = mail.RetryInbucket(5, func() error {
@@ -3428,9 +3421,11 @@ func TestInviteUsersToTeam(t *testing.T) {
 	require.NoError(t, err)
 	nameFormat := *th.App.Config().TeamSettings.TeammateNameDisplay
 	expectedSubject := i18n.T("api.templates.invite_subject",
-		map[string]any{"SenderName": th.SystemAdminUser.GetDisplayName(nameFormat),
+		map[string]any{
+			"SenderName":      th.SystemAdminUser.GetDisplayName(nameFormat),
 			"TeamDisplayName": th.BasicTeam.DisplayName,
-			"SiteName":        th.App.ClientConfig()["SiteName"]})
+			"SiteName":        th.App.ClientConfig()["SiteName"],
+		})
 	checkEmail(t, expectedSubject)
 
 	// Test the invite to team and channel
@@ -3441,10 +3436,12 @@ func TestInviteUsersToTeam(t *testing.T) {
 	_, _, err = th.SystemAdminClient.InviteUsersToTeamAndChannelsGracefully(context.Background(), th.BasicTeam.Id, []string{user1, user2}, []string{th.BasicChannel.Id}, "")
 	require.NoError(t, err)
 	expectedSubject = i18n.T("api.templates.invite_team_and_channel_subject",
-		map[string]any{"SenderName": th.SystemAdminUser.GetDisplayName(nameFormat),
+		map[string]any{
+			"SenderName":      th.SystemAdminUser.GetDisplayName(nameFormat),
 			"TeamDisplayName": th.BasicTeam.DisplayName,
 			"ChannelName":     th.BasicChannel.DisplayName,
-			"SiteName":        th.App.ClientConfig()["SiteName"]})
+			"SiteName":        th.App.ClientConfig()["SiteName"],
+		})
 	checkEmail(t, expectedSubject)
 
 	err = mail.DeleteMailBox(user1)
@@ -3454,9 +3451,11 @@ func TestInviteUsersToTeam(t *testing.T) {
 	_, err = th.LocalClient.InviteUsersToTeam(context.Background(), th.BasicTeam.Id, emailList)
 	require.NoError(t, err)
 	expectedSubject = i18n.T("api.templates.invite_subject",
-		map[string]any{"SenderName": "Administrator",
+		map[string]any{
+			"SenderName":      "Administrator",
 			"TeamDisplayName": th.BasicTeam.DisplayName,
-			"SiteName":        th.App.ClientConfig()["SiteName"]})
+			"SiteName":        th.App.ClientConfig()["SiteName"],
+		})
 	checkEmail(t, expectedSubject)
 
 	// Test the invite local to team and channel
@@ -3467,10 +3466,12 @@ func TestInviteUsersToTeam(t *testing.T) {
 	_, _, err = th.LocalClient.InviteUsersToTeamAndChannelsGracefully(context.Background(), th.BasicTeam.Id, []string{user1, user2}, []string{th.BasicChannel.Id}, "")
 	require.NoError(t, err)
 	expectedSubject = i18n.T("api.templates.invite_team_and_channel_subject",
-		map[string]any{"SenderName": "Administrator",
+		map[string]any{
+			"SenderName":      "Administrator",
 			"TeamDisplayName": th.BasicTeam.DisplayName,
 			"ChannelName":     th.BasicChannel.DisplayName,
-			"SiteName":        th.App.ClientConfig()["SiteName"]})
+			"SiteName":        th.App.ClientConfig()["SiteName"],
+		})
 	checkEmail(t, expectedSubject)
 
 	th.App.UpdateConfig(func(cfg *model.Config) { *cfg.TeamSettings.RestrictCreationToDomains = "@global.com,@common.com" })
@@ -3541,7 +3542,7 @@ func TestInviteGuestsToTeam(t *testing.T) {
 
 	emailList := []string{guest1, guest2}
 
-	//Delete all the messages before check the sample email
+	// Delete all the messages before check the sample email
 	err := mail.DeleteMailBox(guest1)
 	require.NoError(t, err)
 	err = mail.DeleteMailBox(guest2)
@@ -3593,11 +3594,13 @@ func TestInviteGuestsToTeam(t *testing.T) {
 
 	nameFormat := *th.App.Config().TeamSettings.TeammateNameDisplay
 	expectedSubject := i18n.T("api.templates.invite_guest_subject",
-		map[string]any{"SenderName": th.SystemAdminUser.GetDisplayName(nameFormat),
+		map[string]any{
+			"SenderName":      th.SystemAdminUser.GetDisplayName(nameFormat),
 			"TeamDisplayName": th.BasicTeam.DisplayName,
-			"SiteName":        th.App.ClientConfig()["SiteName"]})
+			"SiteName":        th.App.ClientConfig()["SiteName"],
+		})
 
-	//Check if the email was send to the right email address
+	// Check if the email was send to the right email address
 	for _, email := range emailList {
 		var resultsMailbox mail.JSONMessageHeaderInbucket
 		err := mail.RetryInbucket(5, func() error {

server/channels/api4/user_test.go

@@ -770,15 +770,8 @@ func TestCreateUserWebSocketEvent(t *testing.T) {
 		_, _, err := guestClient.Login(context.Background(), guest.Email, guestPassword)
 		require.NoError(t, err)
 
-		guestWSClient, err := th.CreateWebSocketClientWithClient(guestClient)
+		guestWSClient := th.CreateConnectedWebSocketClientWithClient(t, guestClient)
-		require.NoError(t, err)
+		userWSClient := th.CreateConnectedWebSocketClient(t)
-		defer guestWSClient.Close()
-		guestWSClient.Listen()
-
-		userWSClient, err := th.CreateWebSocketClient()
-		require.NoError(t, err)
-		defer userWSClient.Close()
-		userWSClient.Listen()
 
 		user := model.User{Email: th.GenerateTestEmail(), Nickname: "Corey Hulen", Password: "hello1", Username: GenerateTestUsername(), Roles: model.SystemAdminRoleId + " " + model.SystemUserRoleId}
 
@@ -2802,25 +2795,12 @@ func TestUpdateUserActive(t *testing.T) {
 
 		th.App.UpdateConfig(func(cfg *model.Config) { *cfg.TeamSettings.EnableUserDeactivation = true })
 
-		webSocketClient, err := th.CreateWebSocketClient()
+		webSocketClient := th.CreateConnectedWebSocketClient(t)
-		assert.NoError(t, err)
-		defer webSocketClient.Close()
 
-		webSocketClient.Listen()
-
-		time.Sleep(300 * time.Millisecond)
 		resp := <-webSocketClient.ResponseChannel
 		require.Equal(t, model.StatusOk, resp.Status)
 
-		adminWebSocketClient, err := th.CreateWebSocketSystemAdminClient()
+		adminWebSocketClient := th.CreateConnectedWebSocketClientWithClient(t, th.SystemAdminClient)
-		assert.NoError(t, err)
-		defer adminWebSocketClient.Close()
-
-		adminWebSocketClient.Listen()
-
-		time.Sleep(300 * time.Millisecond)
-		resp = <-adminWebSocketClient.ResponseChannel
-		require.Equal(t, model.StatusOk, resp.Status)
 
 		// Verify that both admins and regular users see the email when privacy settings allow same,
 		// and confirm event is fired for SystemAdmin and Local mode
@@ -6288,27 +6268,13 @@ func TestDemoteUserToGuest(t *testing.T) {
 	}, "demote a user to guest")
 
 	t.Run("websocket update user event", func(t *testing.T) {
-		webSocketClient, err := th.CreateWebSocketClient()
+		webSocketClient := th.CreateConnectedWebSocketClient(t)
-		assert.NoError(t, err)
-		defer webSocketClient.Close()
-
-		webSocketClient.Listen()
-
-		time.Sleep(300 * time.Millisecond)
 		resp := <-webSocketClient.ResponseChannel
 		require.Equal(t, model.StatusOk, resp.Status)
 
-		adminWebSocketClient, err := th.CreateWebSocketSystemAdminClient()
+		adminWebSocketClient := th.CreateConnectedWebSocketClientWithClient(t, th.SystemAdminClient)
-		assert.NoError(t, err)
-		defer adminWebSocketClient.Close()
 
-		adminWebSocketClient.Listen()
+		_, _, err := th.SystemAdminClient.GetUser(context.Background(), user.Id, "")
-
-		time.Sleep(300 * time.Millisecond)
-		resp = <-adminWebSocketClient.ResponseChannel
-		require.Equal(t, model.StatusOk, resp.Status)
-
-		_, _, err = th.SystemAdminClient.GetUser(context.Background(), user.Id, "")
 		require.NoError(t, err)
 		_, err = th.SystemAdminClient.DemoteUserToGuest(context.Background(), user.Id)
 		require.NoError(t, err)
@@ -6358,28 +6324,14 @@ func TestPromoteGuestToUser(t *testing.T) {
 	}, "promote a guest to user")
 
 	t.Run("websocket update user event", func(t *testing.T) {
-		t.Skip("https://mattermost.atlassian.net/browse/MM-61736")
+		webSocketClient := th.CreateConnectedWebSocketClient(t)
-		webSocketClient, err := th.CreateWebSocketClient()
-		assert.NoError(t, err)
-		defer webSocketClient.Close()
 
-		webSocketClient.Listen()
-
-		time.Sleep(300 * time.Millisecond)
 		resp := <-webSocketClient.ResponseChannel
 		require.Equal(t, model.StatusOk, resp.Status)
 
-		adminWebSocketClient, err := th.CreateWebSocketSystemAdminClient()
+		adminWebSocketClient := th.CreateConnectedWebSocketClientWithClient(t, th.SystemAdminClient)
-		assert.NoError(t, err)
-		defer adminWebSocketClient.Close()
 
-		adminWebSocketClient.Listen()
+		_, _, err := th.SystemAdminClient.GetUser(context.Background(), user.Id, "")
-
-		time.Sleep(300 * time.Millisecond)
-		resp = <-adminWebSocketClient.ResponseChannel
-		require.Equal(t, model.StatusOk, resp.Status)
-
-		_, _, err = th.SystemAdminClient.GetUser(context.Background(), user.Id, "")
 		require.NoError(t, err)
 		_, err = th.SystemAdminClient.PromoteGuestToUser(context.Background(), user.Id)
 		require.NoError(t, err)
@@ -6584,17 +6536,12 @@ func TestPublishUserTyping(t *testing.T) {
 	})
 
 	t.Run("should send typing event via websocket when triggering a typing event for a user with a common channel", func(t *testing.T) {
-		webSocketClient, err := th.CreateWebSocketClient()
+		webSocketClient := th.CreateConnectedWebSocketClient(t)
-		assert.NoError(t, err)
-		defer webSocketClient.Close()
 
-		webSocketClient.Listen()
-
-		time.Sleep(300 * time.Millisecond)
 		wsResp := <-webSocketClient.ResponseChannel
 		require.Equal(t, model.StatusOk, wsResp.Status)
 
-		_, err = th.SystemAdminClient.PublishUserTyping(context.Background(), th.BasicUser2.Id, tr)
+		_, err := th.SystemAdminClient.PublishUserTyping(context.Background(), th.BasicUser2.Id, tr)
 		require.NoError(t, err)
 
 		assertExpectedWebsocketEvent(t, webSocketClient, model.WebsocketEventTyping, func(resp *model.WebSocketEvent) {
@@ -7240,10 +7187,7 @@ func TestThreadSocketEvents(t *testing.T) {
 		*cfg.ServiceSettings.CollapsedThreads = model.CollapsedThreadsDefaultOn
 	})
 
-	userWSClient, err := th.CreateWebSocketClient()
+	userWSClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	defer userWSClient.Close()
-	userWSClient.Listen()
 
 	client := th.Client
 
@@ -8280,19 +8224,13 @@ func TestUserUpdateEvents(t *testing.T) {
 
 	client1 := th.CreateClient()
 	th.LoginBasicWithClient(client1)
-	WebSocketClient, err := th.CreateWebSocketClientWithClient(client1)
+	WebSocketClient := th.CreateConnectedWebSocketClientWithClient(t, client1)
-	require.NoError(t, err)
-	defer WebSocketClient.Close()
-	WebSocketClient.Listen()
 	resp := <-WebSocketClient.ResponseChannel
 	require.Equal(t, resp.Status, model.StatusOk)
 
 	client2 := th.CreateClient()
 	th.LoginBasic2WithClient(client2)
-	WebSocketClient2, err := th.CreateWebSocketClientWithClient(client2)
+	WebSocketClient2 := th.CreateConnectedWebSocketClientWithClient(t, client2)
-	require.NoError(t, err)
-	defer WebSocketClient2.Close()
-	WebSocketClient2.Listen()
 	resp = <-WebSocketClient2.ResponseChannel
 	require.Equal(t, resp.Status, model.StatusOk)
 

server/channels/api4/websocket_test.go

@@ -33,11 +33,7 @@ func TestWebSocketEvent(t *testing.T) {
 	th := Setup(t).InitBasic()
 	defer th.TearDown()
 
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	defer WebSocketClient.Close()
-
-	WebSocketClient.Listen()
 
 	resp := <-WebSocketClient.ResponseChannel
 	require.Equal(t, resp.Status, model.StatusOk, "should have responded OK to authentication challenge")
@@ -243,22 +239,18 @@ func TestWebSocketSendBinary(t *testing.T) {
 
 	client := th.CreateClient()
 	th.LoginBasicWithClient(client)
-	WebSocketClient, err := th.CreateWebSocketClientWithClient(client)
+	WebSocketClient := th.CreateConnectedWebSocketClientWithClient(t, client)
-	require.NoError(t, err)
-	defer WebSocketClient.Close()
-	WebSocketClient.Listen()
 	resp := <-WebSocketClient.ResponseChannel
 	require.Equal(t, resp.Status, model.StatusOk)
 
 	client2 := th.CreateClient()
 	th.LoginBasic2WithClient(client2)
-	WebSocketClient2, err := th.CreateWebSocketClientWithClient(client2)
+	_ = th.CreateConnectedWebSocketClientWithClient(t, client2)
-	require.NoError(t, err)
-	defer WebSocketClient2.Close()
 
-	time.Sleep(1000 * time.Millisecond)
+	// Wait for statuses to be updated
+	time.Sleep(time.Second)
 
-	err = WebSocketClient.SendBinaryMessage("get_statuses", nil)
+	err := WebSocketClient.SendBinaryMessage("get_statuses", nil)
 	require.NoError(t, err)
 	resp = <-WebSocketClient.ResponseChannel
 	require.Nil(t, resp.Error, resp.Error)
@@ -285,10 +277,7 @@ func TestWebSocketStatuses(t *testing.T) {
 	defer th.TearDown()
 
 	client := th.Client
-	WebSocketClient, err := th.CreateWebSocketClient()
+	WebSocketClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	defer WebSocketClient.Close()
-	WebSocketClient.Listen()
 
 	resp := <-WebSocketClient.ResponseChannel
 	require.Equal(t, resp.Status, model.StatusOk, "should have responded OK to authentication challenge")
@@ -315,10 +304,10 @@ func TestWebSocketStatuses(t *testing.T) {
 
 	th.LoginBasic2()
 
-	WebSocketClient2, err := th.CreateWebSocketClient()
+	WebSocketClient2 := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
 
-	time.Sleep(1000 * time.Millisecond)
+	// Wait for statuses to be updated
+	time.Sleep(time.Second)
 
 	WebSocketClient.GetStatuses()
 	resp = <-WebSocketClient.ResponseChannel
@@ -426,10 +415,7 @@ func TestWebSocketPresence(t *testing.T) {
 	th := Setup(t).InitBasic()
 	defer th.TearDown()
 
-	wsClient, err := th.CreateWebSocketClient()
+	wsClient := th.CreateConnectedWebSocketClient(t)
-	require.NoError(t, err)
-	defer wsClient.Close()
-	wsClient.Listen()
 
 	resp := <-wsClient.ResponseChannel
 	require.Equal(t, resp.Status, model.StatusOk, "should have responded OK to authentication challenge")