From d04b69c9cb21b0b8b15cd3fc50e034e335e95183 Mon Sep 17 00:00:00 2001
From: Claudio Costa <cstcld91@gmail.com>
Date: Wed, 23 Feb 2022 16:01:23 +0100
Subject: [PATCH] Fix racy TestReliableWebSocketSend (#19620)

---
 app/helper_test.go  | 24 ++++++++++++++++++------
 app/options.go      |  8 ++++++++
 app/web_hub_test.go | 33 +++++++++++++++++++++++----------
 3 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/app/helper_test.go b/app/helper_test.go
index e3f7d5c12f..90d6893cd4 100644
--- a/app/helper_test.go
+++ b/app/helper_test.go
@@ -18,6 +18,7 @@ import (
 
 	"github.com/mattermost/mattermost-server/v6/app/request"
 	"github.com/mattermost/mattermost-server/v6/config"
+	"github.com/mattermost/mattermost-server/v6/einterfaces"
 	"github.com/mattermost/mattermost-server/v6/model"
 	"github.com/mattermost/mattermost-server/v6/plugin"
 	"github.com/mattermost/mattermost-server/v6/shared/mlog"
@@ -46,7 +47,7 @@ type TestHelper struct {
 	tempWorkspace string
 }
 
-func setupTestHelper(dbStore store.Store, enterprise bool, includeCacheLayer bool, tb testing.TB) *TestHelper {
+func setupTestHelper(dbStore store.Store, enterprise bool, includeCacheLayer bool, options []Option, tb testing.TB) *TestHelper {
 	tempWorkspace, err := ioutil.TempDir("", "apptest")
 	if err != nil {
 		panic(err)
@@ -65,7 +66,6 @@ func setupTestHelper(dbStore store.Store, enterprise bool, includeCacheLayer boo
 
 	buffer := &mlog.Buffer{}
 
-	var options []Option
 	options = append(options, ConfigStore(configStore))
 	if includeCacheLayer {
 		// Adds the cache layer to the test store
@@ -156,7 +156,7 @@ func Setup(tb testing.TB) *TestHelper {
 	dbStore.MarkSystemRanUnitTests()
 	mainHelper.PreloadMigrations()
 
-	return setupTestHelper(dbStore, false, true, tb)
+	return setupTestHelper(dbStore, false, true, nil, tb)
 }
 
 func SetupWithoutPreloadMigrations(tb testing.TB) *TestHelper {
@@ -167,12 +167,12 @@ func SetupWithoutPreloadMigrations(tb testing.TB) *TestHelper {
 	dbStore.DropAllTables()
 	dbStore.MarkSystemRanUnitTests()
 
-	return setupTestHelper(dbStore, false, true, tb)
+	return setupTestHelper(dbStore, false, true, nil, tb)
 }
 
 func SetupWithStoreMock(tb testing.TB) *TestHelper {
 	mockStore := testlib.GetMockStoreForSetupFunctions()
-	th := setupTestHelper(mockStore, false, false, tb)
+	th := setupTestHelper(mockStore, false, false, nil, tb)
 	statusMock := mocks.StatusStore{}
 	statusMock.On("UpdateExpiredDNDStatuses").Return([]*model.Status{}, nil)
 	statusMock.On("Get", "user1").Return(&model.Status{UserId: "user1", Status: model.StatusOnline}, nil)
@@ -187,7 +187,7 @@ func SetupWithStoreMock(tb testing.TB) *TestHelper {
 
 func SetupEnterpriseWithStoreMock(tb testing.TB) *TestHelper {
 	mockStore := testlib.GetMockStoreForSetupFunctions()
-	th := setupTestHelper(mockStore, true, false, tb)
+	th := setupTestHelper(mockStore, true, false, nil, tb)
 	statusMock := mocks.StatusStore{}
 	statusMock.On("UpdateExpiredDNDStatuses").Return([]*model.Status{}, nil)
 	statusMock.On("Get", "user1").Return(&model.Status{UserId: "user1", Status: model.StatusOnline}, nil)
@@ -200,6 +200,18 @@ func SetupEnterpriseWithStoreMock(tb testing.TB) *TestHelper {
 	return th
 }
 
+func SetupWithClusterMock(tb testing.TB, cluster einterfaces.ClusterInterface) *TestHelper {
+	if testing.Short() {
+		tb.SkipNow()
+	}
+	dbStore := mainHelper.GetStore()
+	dbStore.DropAllTables()
+	dbStore.MarkSystemRanUnitTests()
+	mainHelper.PreloadMigrations()
+
+	return setupTestHelper(dbStore, true, true, []Option{setCluster(cluster)}, tb)
+}
+
 var initBasicOnce sync.Once
 var userCache struct {
 	SystemAdminUser *model.User
diff --git a/app/options.go b/app/options.go
index a88edb51f3..5da342136e 100644
--- a/app/options.go
+++ b/app/options.go
@@ -7,6 +7,7 @@ import (
 	"github.com/pkg/errors"
 
 	"github.com/mattermost/mattermost-server/v6/config"
+	"github.com/mattermost/mattermost-server/v6/einterfaces"
 	"github.com/mattermost/mattermost-server/v6/model"
 	"github.com/mattermost/mattermost-server/v6/shared/mlog"
 	"github.com/mattermost/mattermost-server/v6/store"
@@ -111,3 +112,10 @@ func ServerConnector(ch *Channels) AppOption {
 		a.ch = ch
 	}
 }
+
+func setCluster(cluster einterfaces.ClusterInterface) Option {
+	return func(s *Server) error {
+		s.Cluster = cluster
+		return nil
+	}
+}
diff --git a/app/web_hub_test.go b/app/web_hub_test.go
index e7ef4f2b87..423a800053 100644
--- a/app/web_hub_test.go
+++ b/app/web_hub_test.go
@@ -338,24 +338,37 @@ func TestHubConnIndexInactive(t *testing.T) {
 }
 
 func TestReliableWebSocketSend(t *testing.T) {
-	th := Setup(t)
+	testCluster := &testlib.FakeClusterInterface{}
+
+	th := SetupWithClusterMock(t, testCluster)
 	defer th.TearDown()
 
-	testCluster := &testlib.FakeClusterInterface{}
-	th.Server.Cluster = testCluster
-
-	ev := model.NewWebSocketEvent("test_reliable_event", "", "", "", nil)
+	ev := model.NewWebSocketEvent("test_unreliable_event", "", "", "", nil)
 	ev = ev.SetBroadcast(&model.WebsocketBroadcast{})
 	th.App.Publish(ev)
-	ev = ev.SetBroadcast(&model.WebsocketBroadcast{
+	ev2 := model.NewWebSocketEvent("test_reliable_event", "", "", "", nil)
+	ev2 = ev2.SetBroadcast(&model.WebsocketBroadcast{
 		ReliableClusterSend: true,
 	})
-	th.App.Publish(ev)
+	th.App.Publish(ev2)
 
 	messages := testCluster.GetMessages()
-	require.Len(t, messages, 2)
-	require.Equal(t, model.ClusterSendBestEffort, messages[0].SendType)
-	require.Equal(t, model.ClusterSendReliable, messages[1].SendType)
+
+	evJSON, err := ev.ToJSON()
+	require.NoError(t, err)
+	ev2JSON, err := ev2.ToJSON()
+	require.NoError(t, err)
+
+	require.Contains(t, messages, &model.ClusterMessage{
+		Event:    model.ClusterEventPublish,
+		Data:     evJSON,
+		SendType: model.ClusterSendBestEffort,
+	})
+	require.Contains(t, messages, &model.ClusterMessage{
+		Event:    model.ClusterEventPublish,
+		Data:     ev2JSON,
+		SendType: model.ClusterSendReliable,
+	})
 }
 
 func TestHubIsRegistered(t *testing.T) {