diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
index f1895a64895..21920c44da7 100644
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -2443,7 +2443,7 @@ def fast_on_a(lock):
 
     assert "Unexpected worker completed task" in s_logs.getvalue()
 
-    while not len(s.tasks[f2.key].who_has) == 2:
+    while a.address not in {w.address for w in s.tasks[f2.key].who_has}:
         await asyncio.sleep(0.001)
 
     # Ensure that all keys have been properly registered and will also be
@@ -2510,7 +2510,7 @@ def fast_on_a(lock):
 
     assert "Unexpected worker completed task" in s_logs.getvalue()
 
-    while not len(s.tasks[f2.key].who_has) == 2:
+    while a.address not in {w.address for w in s.tasks[f2.key].who_has}:
         await asyncio.sleep(0.001)
 
     del f1, f2, f3