Fix flaky test_worker_reconnects_mid_compute (#5299)

distributed/tests/test_worker.py

@@ -2409,11 +2409,6 @@ async def test_hold_on_to_replicas(c, s, *workers):
         await asyncio.sleep(0.01)
 
 
-@pytest.mark.flaky(
-    condition=WINDOWS and sys.version_info[:2] == (3, 8),
-    reruns=20,
-    reruns_delay=5,
-)
 @gen_cluster(client=True)
 async def test_worker_reconnects_mid_compute(c, s, a, b):
     """Ensure that, if a worker disconnects while computing a result, the scheduler will
@@ -2468,8 +2463,11 @@ def fast_on_a(lock):
 
     assert "Unexpected worker completed task" in s_logs.getvalue()
 
-    while a.address not in {w.address for w in s.tasks[f2.key].who_has}:
-        await asyncio.sleep(0.001)
+    # Ensure that all in-memory tasks on A have been restored on the
+    # scheduler after reconnect
+    for ts in a.tasks.values():
+        if ts.state == "memory":
+            assert a.address in {ws.address for ws in s.tasks[ts.key].who_has}
 
     # Ensure that all keys have been properly registered and will also be
     # cleaned up nicely.
@@ -2479,12 +2477,7 @@ def fast_on_a(lock):
         await asyncio.sleep(0.001)
 
 
-@pytest.mark.flaky(
-    condition=WINDOWS and sys.version_info[:2] == (3, 8),
-    reruns=20,
-    reruns_delay=5,
-)
-@gen_cluster(client=True)
+@gen_cluster(client=True, timeout=5)
 async def test_worker_reconnects_mid_compute_multiple_states_on_scheduler(c, s, a, b):
     """
     Ensure that a reconnecting worker does not break the scheduler regardless of
@@ -2540,8 +2533,11 @@ def fast_on_a(lock):
 
     assert "Unexpected worker completed task" in s_logs.getvalue()
 
-    while a.address not in {w.address for w in s.tasks[f2.key].who_has}:
-        await asyncio.sleep(0.001)
+    # Ensure that all in-memory tasks on A have been restored on the
+    # scheduler after reconnect
+    for ts in a.tasks.values():
+        if ts.state == "memory":
+            assert a.address in {ws.address for ws in s.tasks[ts.key].who_has}
 
     del f1, f2, f3
     while any(w.tasks for w in [a, b]):