From 8890f41491aad09afbd661862eec1e48bd7314bd Mon Sep 17 00:00:00 2001
From: Ibrahim Serdar Acikgoz <serdaracikgoz86@gmail.com>
Date: Thu, 4 Jun 2020 12:52:40 +0300
Subject: [PATCH] api4/post_test: fix a flaky test (#14741)

---
 api4/post_test.go | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/api4/post_test.go b/api4/post_test.go
index 5672ab82e3..dd5283d9fe 100644
--- a/api4/post_test.go
+++ b/api4/post_test.go
@@ -602,9 +602,8 @@ func TestCreatePostCheckOnlineStatus(t *testing.T) {
 	wsClient, err := th.CreateWebSocketClientWithClient(cli)
 	require.Nil(t, err)
 	defer func() {
-		wsClient.Close()
 		wg.Wait()
-
+		wsClient.Close()
 	}()
 
 	wsClient.Listen()
@@ -612,15 +611,23 @@ func TestCreatePostCheckOnlineStatus(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		timeout := time.After(2 * time.Second)
+		cnt := true
 		i := 0
-		for ev := range wsClient.EventChannel {
-			if ev.EventType() == model.WEBSOCKET_EVENT_POSTED {
-				if i == 0 {
-					assert.False(t, ev.GetData()["set_online"].(bool))
-				} else {
-					assert.True(t, ev.GetData()["set_online"].(bool))
+		for cnt {
+			select {
+			case ev := <-wsClient.EventChannel:
+				if ev.EventType() == model.WEBSOCKET_EVENT_POSTED {
+					if i == 0 {
+						assert.False(t, ev.GetData()["set_online"].(bool))
+					} else {
+						assert.True(t, ev.GetData()["set_online"].(bool))
+					}
+					i++
 				}
-				i++
+				cnt = i != 2
+			case <-timeout:
+				cnt = false
 			}
 		}
 		assert.Equal(t, 2, i, "unexpected number of posted events")