From d2e2646eb9827bb883d382267d1c13fe0e61d153 Mon Sep 17 00:00:00 2001
From: Alexey Kiselev <alexey.kiselev@gmail.com>
Date: Thu, 19 Dec 2024 16:54:17 +0400
Subject: [PATCH] Better test workflow. Better wait group naming.

---
 pkg/networking/session_test.go | 59 ++++++++++++++++++++--------------
 1 file changed, 34 insertions(+), 25 deletions(-)

diff --git a/pkg/networking/session_test.go b/pkg/networking/session_test.go
index da83bcaff..fddc2c38c 100644
--- a/pkg/networking/session_test.go
+++ b/pkg/networking/session_test.go
@@ -170,42 +170,51 @@ func TestSessionTimeoutOnMessage(t *testing.T) {
 
 	serverHandler.On("OnClose", serverSession).Return()
 
-	var serverWG sync.WaitGroup
-	var clientWG sync.WaitGroup
-	serverWG.Add(1)
-	clientWG.Add(1)
-	go func() {
-		sc1 := serverHandler.On("OnHandshake", serverSession, &textHandshake{v: "hello"}).Once().Return()
-		sc1.Run(func(_ mock.Arguments) {
-			n, wErr := serverSession.Write([]byte("hello"))
-			require.NoError(t, wErr)
-			assert.Equal(t, 5, n)
-			serverWG.Done()
-		})
-		serverWG.Wait() // Wait for finishing handshake before closing the pipe.
+	clientWG := new(sync.WaitGroup)
+	clientWG.Add(1) // Wait for client to send Handshake to server.
 
-		// Lock pipe after replying with the handshake from server.
-		pc.writeBlocker.Lock()
-		clientWG.Done() // Signal that pipe is locked.
-	}()
+	serverWG := new(sync.WaitGroup)
+	serverWG.Add(1) // Wait for server to reply with Handshake to client.
 
-	serverHandler.On("OnClose", serverSession).Return()
-	clientHandler.On("OnClose", clientSession).Return()
+	pipeWG := new(sync.WaitGroup)
+	pipeWG.Add(1) // Wait for pipe to be locked.
 
-	// Send handshake to server.
-	n, err := clientSession.Write([]byte("hello"))
-	require.NoError(t, err)
-	assert.Equal(t, 5, n)
+	testWG := new(sync.WaitGroup)
+	testWG.Add(1) // Wait for client fail by timeout.
+
+	serverHandler.On("OnClose", serverSession).Return()
+	sc1 := serverHandler.On("OnHandshake", serverSession, &textHandshake{v: "hello"}).Once().Return()
+	sc1.Run(func(_ mock.Arguments) {
+		clientWG.Wait() // Wait for client to send handshake, start replying with Handshake only after that.
+		n, wErr := serverSession.Write([]byte("hello"))
+		require.NoError(t, wErr)
+		assert.Equal(t, 5, n)
+		serverWG.Done()
+	})
 
+	clientHandler.On("OnClose", clientSession).Return()
 	cs1 := clientHandler.On("OnHandshake", clientSession, &textHandshake{v: "hello"}).Once().Return()
 	cs1.Run(func(_ mock.Arguments) {
-		clientWG.Wait() // Wait for pipe to be locked.
+		pipeWG.Wait() // Wait for pipe to be locked.
 		// On receiving handshake from server, send the message back to server.
 		_, msgErr := clientSession.Write(encodeMessage("Hello session"))
 		require.ErrorIs(t, msgErr, networking.ErrConnectionWriteTimeout)
+		testWG.Done()
 	})
 
-	time.Sleep(1 * time.Second) // Let timeout occur.
+	go func() {
+		serverWG.Wait()        // Wait for finishing handshake before closing the pipe.
+		pc.writeBlocker.Lock() // Lock pipe after replying with the handshake from server.
+		pipeWG.Done()          // Signal that pipe is locked.
+	}()
+
+	// Send handshake to server.
+	n, err := clientSession.Write([]byte("hello"))
+	require.NoError(t, err)
+	assert.Equal(t, 5, n)
+	clientWG.Done() // Signal that handshake was sent to server.
+
+	testWG.Wait()
 
 	err = serverSession.Close()
 	assert.NoError(t, err) // Expect no error on the server side.