diff --git a/client/clientimpl_test.go b/client/clientimpl_test.go
index 97927fc8..522c4186 100644
--- a/client/clientimpl_test.go
+++ b/client/clientimpl_test.go
@@ -626,9 +626,9 @@ func TestAgentIdentification(t *testing.T) {
 		newInstanceUid := ulid.MustNew(
 			ulid.Timestamp(time.Now()), ulid.Monotonic(rand.New(rand.NewSource(0)), 0),
 		)
-		var rcvAgentInstanceUid atomic.Value
+		rcvAgentInstanceUids := make(chan string, 10)
 		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-			rcvAgentInstanceUid.Store(msg.InstanceUid)
+			rcvAgentInstanceUids <- msg.InstanceUid
 			return &protobufs.ServerToAgent{
 				InstanceUid: msg.InstanceUid,
 				AgentIdentification: &protobufs.AgentIdentification{
@@ -638,7 +638,17 @@ func TestAgentIdentification(t *testing.T) {
 		}
 
 		// Start a client.
-		settings := types.StartSettings{}
+		hasRecvServerResponse := make(chan struct{}, 1)
+		settings := types.StartSettings{
+			Callbacks: types.CallbacksStruct{
+				OnMessageFunc: func(_ context.Context, _ *types.MessageData) {
+					select {
+					case hasRecvServerResponse <- struct{}{}:
+					default:
+					}
+				},
+			},
+		}
 		settings.OpAMPServerURL = "ws://" + srv.Endpoint
 		prepareClient(t, &settings, client)
 
@@ -649,14 +659,18 @@ func TestAgentIdentification(t *testing.T) {
 		eventually(
 			t,
 			func() bool {
-				instanceUid, ok := rcvAgentInstanceUid.Load().(string)
-				if !ok {
-					return false
+				select {
+				case instanceUid := <-rcvAgentInstanceUids:
+					return instanceUid == oldInstanceUid
+				default:
 				}
-				return instanceUid == oldInstanceUid
+				return false
 			},
 		)
 
+		// Wait for the server response.
+		<-hasRecvServerResponse
+
 		// Send a dummy message
 		_ = client.SetAgentDescription(createAgentDescr())
 
@@ -665,11 +679,12 @@ func TestAgentIdentification(t *testing.T) {
 		eventually(
 			t,
 			func() bool {
-				instanceUid, ok := rcvAgentInstanceUid.Load().(string)
-				if !ok {
-					return false
+				select {
+				case instanceUid := <-rcvAgentInstanceUids:
+					return instanceUid == newInstanceUid.String()
+				default:
 				}
-				return instanceUid == newInstanceUid.String()
+				return false
 			},
 		)