deflake flaky module watcher tests (#9367)

Author: akshayka

tests/_runtime/reload/test_module_watcher.py

@@ -151,7 +151,16 @@ def foo():
     )
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(INTERVAL * 3)
+    retries = 0
+    while retries < 10:
+        await asyncio.sleep(INTERVAL)
+        retries += 1
+        if (
+            k.graph.cells[er_1.cell_id].stale
+            and k.graph.cells[er_2.cell_id].stale
+        ):
+            break
+
     assert k.graph.cells[er_1.cell_id].stale
     assert k.graph.cells[er_2.cell_id].stale
     assert not k.graph.cells[er_3.cell_id].stale
@@ -162,6 +171,7 @@ def foo():
     assert k.globals["x"] == 2
 
 
+@pytest.mark.flaky(reruns=3)
 async def test_reload_nested_module_function(
     tmp_path: pathlib.Path,
     py_modname: str,
@@ -206,7 +216,16 @@ def foo():
     update_file(nested_module, "func = lambda : 2")
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(INTERVAL * 3)
+    retries = 0
+    while retries < 10:
+        await asyncio.sleep(INTERVAL)
+        retries += 1
+        if (
+            k.graph.cells[er_1.cell_id].stale
+            and k.graph.cells[er_2.cell_id].stale
+        ):
+            break
+
     assert k.graph.cells[er_1.cell_id].stale
     assert k.graph.cells[er_2.cell_id].stale
     assert not k.graph.cells[er_3.cell_id].stale
@@ -217,6 +236,7 @@ def foo():
     assert k.globals["x"] == 2
 
 
+@pytest.mark.flaky(reruns=3)
 async def test_reload_nested_module_import_module(
     tmp_path: pathlib.Path,
     py_modname: str,
@@ -261,7 +281,16 @@ def foo():
     update_file(nested_module, "func = lambda : 2")
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(INTERVAL * 3)
+    retries = 0
+    while retries < 10:
+        await asyncio.sleep(INTERVAL)
+        retries += 1
+        if (
+            k.graph.cells[er_1.cell_id].stale
+            and k.graph.cells[er_2.cell_id].stale
+        ):
+            break
+
     assert k.graph.cells[er_1.cell_id].stale
     assert k.graph.cells[er_2.cell_id].stale
     assert not k.graph.cells[er_3.cell_id].stale
@@ -335,6 +364,7 @@ def foo():
     assert k.globals["x"] == 2
 
 
+@pytest.mark.flaky(reruns=3)
 async def test_reload_package(
     tmp_path: pathlib.Path,
     execution_kernel: Kernel,
@@ -368,7 +398,16 @@ async def test_reload_package(
     update_file(nested_module, "func = lambda : 2")
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(INTERVAL * 3)
+    retries = 0
+    while retries < 10:
+        await asyncio.sleep(INTERVAL)
+        retries += 1
+        if (
+            k.graph.cells[er_1.cell_id].stale
+            and k.graph.cells[er_2.cell_id].stale
+        ):
+            break
+
     assert k.graph.cells[er_1.cell_id].stale
     assert k.graph.cells[er_2.cell_id].stale
     assert not k.graph.cells[er_3.cell_id].stale
@@ -441,6 +480,7 @@ def foo():
     assert k.globals["x"] == 2
 
 
+@pytest.mark.flaky(reruns=3)
 async def test_reload_with_modified_cell(
     tmp_path: pathlib.Path,
     py_modname: str,
@@ -479,7 +519,16 @@ def foo():
     )
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(INTERVAL * 3)
+    retries = 0
+    while retries < 10:
+        await asyncio.sleep(INTERVAL)
+        retries += 1
+        if (
+            k.graph.cells[er_1.cell_id].stale
+            and k.graph.cells[er_2.cell_id].stale
+        ):
+            break
+
     assert k.graph.cells[er_1.cell_id].stale
     assert k.graph.cells[er_2.cell_id].stale
     assert not k.graph.cells[er_3.cell_id].stale
@@ -924,6 +973,7 @@ async def test_module_watcher_processes_flag(
 class TestModuleWatcherEdgeCases:
     """Tests for edge cases in module watching"""
 
+    @pytest.mark.flaky(reruns=3)
     async def test_module_watcher_handles_deleted_cell(
         self,
         tmp_path: pathlib.Path,
@@ -969,12 +1019,18 @@ def foo():
         )
 
         # Wait for watcher to pick up change
-        await asyncio.sleep(INTERVAL * 3)
+        retries = 0
+        while retries < 10:
+            await asyncio.sleep(INTERVAL)
+            retries += 1
+            if k.graph.cells[er_1.cell_id].stale:
+                break
 
         # Only er_1 should be stale (er_2 was deleted)
         assert k.graph.cells[er_1.cell_id].stale
         assert er_2.cell_id not in k.graph.cells
 
+    @pytest.mark.flaky(reruns=3)
     async def test_module_watcher_cache_invalidation(
         self,
         tmp_path: pathlib.Path,
@@ -999,7 +1055,13 @@ async def test_module_watcher_cache_invalidation(
         update_file(py_file, "def foo(): return 2")
 
         # Wait for watcher
-        await asyncio.sleep(INTERVAL * 3)
+        retries = 0
+        while retries < 10:
+            await asyncio.sleep(INTERVAL)
+            retries += 1
+            if k.graph.cells[er_1.cell_id].stale:
+                break
+
         assert k.graph.cells[er_1.cell_id].stale
 
         # Run stale cells