From 77c583b9dfc2a9b6a924cfb8973eebd98596ef2d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Wed, 1 Apr 2026 12:30:46 +0200
Subject: [PATCH] [client] Replace `receivedAcks` map with atomic counter and
 improve stream establishment sync in tests

Refactor acknowledgment tracking in tests to use an `atomic.Int32` counter instead of a map. Replace fixed sleep with robust synchronization by waiting on `handlerStarted` signal for stream establishment.
---
 flow/client/client_test.go | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/flow/client/client_test.go b/flow/client/client_test.go
index 148b193c5..55157acbc 100644
--- a/flow/client/client_test.go
+++ b/flow/client/client_test.go
@@ -197,16 +197,13 @@ func TestReceive(t *testing.T) {
 		assert.NoError(t, err, "failed to close flow")
 	})
 
-	receivedAcks := make(map[string]bool)
+	var ackCount atomic.Int32
 	receiveDone := make(chan struct{})
 
 	go func() {
 		err := client.Receive(ctx, 1*time.Second, func(msg *proto.FlowEventAck) error {
 			if !msg.IsInitiator && len(msg.EventId) > 0 {
-				id := string(msg.EventId)
-				receivedAcks[id] = true
-
-				if len(receivedAcks) >= 3 {
+				if ackCount.Add(1) >= 3 {
 					close(receiveDone)
 				}
 			}
@@ -217,7 +214,11 @@ func TestReceive(t *testing.T) {
 		}
 	}()
 
-	time.Sleep(500 * time.Millisecond)
+	select {
+	case <-server.handlerStarted:
+	case <-time.After(3 * time.Second):
+		t.Fatal("timeout waiting for stream to be established")
+	}
 
 	for i := 0; i < 3; i++ {
 		eventID := uuid.New().String()
@@ -240,7 +241,7 @@ func TestReceive(t *testing.T) {
 		t.Fatal("timeout waiting for acks to be processed")
 	}
 
-	assert.Equal(t, 3, len(receivedAcks))
+	assert.Equal(t, int32(3), ackCount.Load())
 }
 
 func TestReceive_ContextCancellation(t *testing.T) {
@@ -486,7 +487,11 @@ func TestReceive_ProtocolErrorStreamReconnect(t *testing.T) {
 	}()
 
 	// Wait for stream to be established, then send first ack
-	time.Sleep(500 * time.Millisecond)
+	select {
+	case <-server.handlerStarted:
+	case <-time.After(3 * time.Second):
+		t.Fatal("timeout waiting for stream to be established")
+	}
 	server.acks <- &proto.FlowEventAck{EventId: []byte("before-close")}
 
 	select {