From 7b3fffbcc59368cec2562d4ab9d08b34a9e250f5 Mon Sep 17 00:00:00 2001
From: Alexey Kiselev <alexey.kiselev@gmail.com>
Date: Thu, 19 Dec 2024 15:34:32 +0400
Subject: [PATCH] Fixed blinking test TestCloseParentContext. Wait group added
 to wait for client to finish sending handshake. Better wait groups naming.

---
 pkg/networking/session_test.go | 46 +++++++++++++++++++---------------
 1 file changed, 26 insertions(+), 20 deletions(-)

diff --git a/pkg/networking/session_test.go b/pkg/networking/session_test.go
index 49759c03a..da83bcaff 100644
--- a/pkg/networking/session_test.go
+++ b/pkg/networking/session_test.go
@@ -330,7 +330,6 @@ func TestCloseParentContext(t *testing.T) {
 	serverHandler := netmocks.NewMockHandler(t)
 
 	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 
 	clientConn, serverConn := testConnPipe()
 	net := networking.NewNetwork()
@@ -340,44 +339,51 @@ func TestCloseParentContext(t *testing.T) {
 	serverSession, err := net.NewSession(ctx, serverConn, testConfig(t, mockProtocol, serverHandler, "server"))
 	require.NoError(t, err)
 
-	var closeWG sync.WaitGroup
-	closeWG.Add(1)
+	clientWG := new(sync.WaitGroup)
+	clientWG.Add(1) // Wait for client to send Handshake to server.
 
-	var wg sync.WaitGroup
-	wg.Add(2)
+	serverWG := new(sync.WaitGroup)
+	serverWG.Add(1) // Wait for server to send Handshake to client, after that we will close the parent context.
+
+	testWG := new(sync.WaitGroup)
+	testWG.Add(2) // Wait for both client and server to finish.
 
 	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"))
 		assert.NoError(t, wErr)
 		assert.Equal(t, 5, n)
 		go func() {
-			closeWG.Wait() // Wait for client to receive server handshake.
-			cancel()       // Close parent context.
-			wg.Done()
+			serverWG.Wait() // Wait for client to receive server handshake.
+			cancel()        // Close parent context.
+			testWG.Done()
 		}()
 	})
 
 	clientHandler.On("OnClose", clientSession).Return()
 
-	// Send handshake to server.
-	n, err := clientSession.Write([]byte("hello"))
-	require.NoError(t, err)
-	assert.Equal(t, 5, n)
-
 	cs1 := clientHandler.On("OnHandshake", clientSession, &textHandshake{v: "hello"}).Once().Return()
 	cs1.Run(func(_ mock.Arguments) {
 		// On receiving handshake from server, signal to close the server.
-		closeWG.Done()
-		// Try to send message to server, but it will fail because server is already closed.
-		time.Sleep(10 * time.Millisecond) // Wait for server to close.
-		_, msgErr := clientSession.Write(encodeMessage("Hello session"))
-		require.ErrorIs(t, msgErr, networking.ErrSessionShutdown)
-		wg.Done()
+		serverWG.Done()
+		go func() {
+			// Try to send message to server, but it will fail because server is already closed.
+			time.Sleep(10 * time.Millisecond) // Wait for server to close.
+			_, msgErr := clientSession.Write(encodeMessage("Hello session"))
+			require.ErrorIs(t, msgErr, networking.ErrSessionShutdown)
+			testWG.Done()
+		}()
 	})
 
-	wg.Wait() // Wait for client to finish.
+	// 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() // Wait for all interactions to finish.
 
 	err = clientSession.Close()
 	assert.NoError(t, err)