diff --git a/tests/unit/cluster/replica-migration.tcl b/tests/unit/cluster/replica-migration.tcl
index 0d232b276..49c31128b 100644
--- a/tests/unit/cluster/replica-migration.tcl
+++ b/tests/unit/cluster/replica-migration.tcl
@@ -339,11 +339,6 @@ start_cluster 4 4 {tags {external:skip cluster} overrides {cluster-node-timeout
         R 3 config set cluster-allow-replica-migration no
         R 7 config set cluster-allow-replica-migration yes
 
-        # Record the current primary node, server 7 will be migrated later.
-        set old_role_response [R 7 role]
-        set old_primary_ip [lindex [R 7 role] 1]
-        set old_primary_port [lindex [R 7 role] 2]
-
         # Move slot 0 from primary 3 to primary 0.
         set addr "[srv 0 host]:[srv 0 port]"
         set myid [R 3 CLUSTER MYID]
@@ -354,26 +349,16 @@ start_cluster 4 4 {tags {external:skip cluster} overrides {cluster-node-timeout
             fail "valkey-cli --cluster rebalance returns non-zero exit code, output below:\n$result"
         }
 
-        # Wait for server 7 role response to change.
+        # Make sure server 3 lost its replica (server 7) and server 7 becomes a replica of primary 0.
         wait_for_condition 1000 50 {
-            [R 7 role] ne $old_role_response
+            [s -3 role] eq {master} &&
+            [s -3 connected_slaves] eq 0 &&
+            [s -7 role] eq {slave} &&
+            [get_my_primary_peer 7] eq $addr
         } else {
             puts "R 3 role: [R 3 role]"
             puts "R 7 role: [R 7 role]"
-            fail "Server 7 role response has not changed"
+            fail "Server 3 and 7 role response has not changed"
         }
-
-        wait_for_cluster_propagation
-        wait_for_cluster_state "ok"
-
-        # Make sure server 3 is still a primary and has no replicas.
-        assert_equal [s -3 role] {master}
-        assert_equal [lindex [R 3 role] 2] {}
-
-        # And server 7 becomes a replica of another primary.
-        set new_primary_ip [lindex [R 7 role] 1]
-        set new_primary_port [lindex [R 7 role] 2]
-        assert_equal [s -7 role] {slave}
-        assert_not_equal "$old_primary_ip:$old_primary_port" "new_primary_ip:new_primary_port"
     }
 } my_slot_allocation cluster_allocate_replicas ;# start_cluster