From a3ce298729a6e5c246b854efd22415eb792bde27 Mon Sep 17 00:00:00 2001
From: Jannis Mattheis <contact@jmattheis.de>
Date: Sun, 6 Aug 2023 12:27:41 +0200
Subject: [PATCH] Fix stream_test flakiness

---
 api/stream/stream_test.go | 58 +++++++++++++++++++++++++++------------
 1 file changed, 41 insertions(+), 17 deletions(-)

diff --git a/api/stream/stream_test.go b/api/stream/stream_test.go
index aaf7cc2..009b10f 100644
--- a/api/stream/stream_test.go
+++ b/api/stream/stream_test.go
@@ -57,8 +57,8 @@ func TestWriteMessageFails(t *testing.T) {
 	wsURL := wsURL(server.URL)
 	user := testClient(t, wsURL)
 
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+	waitForConnectedClients(api, 1)
+
 	clients := clients(api, 1)
 	assert.NotEmpty(t, clients)
 
@@ -87,13 +87,13 @@ func TestWritePingFails(t *testing.T) {
 	user := testClient(t, wsURL)
 	defer user.conn.Close()
 
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+	waitForConnectedClients(api, 1)
+
 	clients := clients(api, 1)
 
 	assert.NotEmpty(t, clients)
 
-	time.Sleep(api.pingPeriod) // waiting for ping
+	time.Sleep(api.pingPeriod + (50 * time.Millisecond)) // waiting for ping
 
 	api.Notify(1, &model.MessageExternal{Message: "HI"})
 	user.expectNoMessage()
@@ -148,8 +148,8 @@ func TestCloseClientOnNotReading(t *testing.T) {
 	assert.Nil(t, err)
 	defer ws.Close()
 
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+	waitForConnectedClients(api, 1)
+
 	assert.NotEmpty(t, clients(api, 1))
 
 	time.Sleep(api.pingPeriod + api.pongTimeout)
@@ -168,8 +168,9 @@ func TestMessageDirectlyAfterConnect(t *testing.T) {
 
 	user := testClient(t, wsURL)
 	defer user.conn.Close()
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+
+	waitForConnectedClients(api, 1)
+
 	api.Notify(1, &model.MessageExternal{Message: "msg"})
 	user.expectMessage(&model.MessageExternal{Message: "msg"})
 }
@@ -185,8 +186,9 @@ func TestDeleteClientShouldCloseConnection(t *testing.T) {
 
 	user := testClient(t, wsURL)
 	defer user.conn.Close()
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+
+	waitForConnectedClients(api, 1)
+
 	api.Notify(1, &model.MessageExternal{Message: "msg"})
 	user.expectMessage(&model.MessageExternal{Message: "msg"})
 
@@ -231,8 +233,7 @@ func TestDeleteMultipleClients(t *testing.T) {
 	defer userThreeAndroid.conn.Close()
 	userThree := []*testingClient{userThreeAndroid}
 
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+	waitForConnectedClients(api, len(userOne)+len(userTwo)+len(userThree))
 
 	api.Notify(1, &model.MessageExternal{ID: 4, Message: "there"})
 	expectMessage(&model.MessageExternal{ID: 4, Message: "there"}, userOne...)
@@ -295,8 +296,7 @@ func TestDeleteUser(t *testing.T) {
 	defer userThreeAndroid.conn.Close()
 	userThree := []*testingClient{userThreeAndroid}
 
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+	waitForConnectedClients(api, len(userOne)+len(userTwo)+len(userThree))
 
 	api.Notify(1, &model.MessageExternal{ID: 4, Message: "there"})
 	expectMessage(&model.MessageExternal{ID: 4, Message: "there"}, userOne...)
@@ -343,6 +343,8 @@ func TestCollectConnectedClientTokens(t *testing.T) {
 	defer userOneConnTwo.conn.Close()
 	userOneConnThree := testClient(t, wsURL)
 	defer userOneConnThree.conn.Close()
+	waitForConnectedClients(api, 3)
+
 	ret := api.CollectConnectedClientTokens()
 	sort.Strings(ret)
 	assert.Equal(t, []string{"1-1", "1-2"}, ret)
@@ -351,6 +353,8 @@ func TestCollectConnectedClientTokens(t *testing.T) {
 	defer userTwoConnOne.conn.Close()
 	userTwoConnTwo := testClient(t, wsURL)
 	defer userTwoConnTwo.conn.Close()
+	waitForConnectedClients(api, 5)
+
 	ret = api.CollectConnectedClientTokens()
 	sort.Strings(ret)
 	assert.Equal(t, []string{"1-1", "1-2", "2-1", "2-2"}, ret)
@@ -388,8 +392,7 @@ func TestMultipleClients(t *testing.T) {
 	defer userThreeAndroid.conn.Close()
 	userThree := []*testingClient{userThreeAndroid}
 
-	// the server may take some time to register the client
-	time.Sleep(100 * time.Millisecond)
+	waitForConnectedClients(api, len(userOne)+len(userTwo)+len(userThree))
 
 	// there should not be messages at the beginning
 	expectNoMessage(userOne...)
@@ -508,6 +511,17 @@ func clients(api *API, user uint) []*client {
 	return api.clients[user]
 }
 
+func countClients(a *API) int {
+	a.lock.RLock()
+	defer a.lock.RUnlock()
+
+	var i int
+	for _, clients := range a.clients {
+		i += len(clients)
+	}
+	return i
+}
+
 func testClient(t *testing.T, url string) *testingClient {
 	client := createClient(t, url)
 	startReading(client)
@@ -594,3 +608,13 @@ func staticUserID() gin.HandlerFunc {
 		auth.RegisterAuthentication(context, nil, 1, "customtoken")
 	}
 }
+
+func waitForConnectedClients(api *API, count int) {
+	for i := 0; i < 10; i++ {
+		if countClients(api) == count {
+			// ok
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}