api4/post_test: fix a flaky test (#14741)

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 {
+		for cnt {
-			if ev.EventType() == model.WEBSOCKET_EVENT_POSTED {
+			select {
-				if i == 0 {
+			case ev := <-wsClient.EventChannel:
-					assert.False(t, ev.GetData()["set_online"].(bool))
+				if ev.EventType() == model.WEBSOCKET_EVENT_POSTED {
-				} else {
+					if i == 0 {
-					assert.True(t, ev.GetData()["set_online"].(bool))
+						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")