fix: faster module watcher tests (#3062)

marimo/_runtime/reload/module_watcher.py

@@ -113,6 +113,12 @@ def _check_modules(
     return stale_modules
 
 
+MODULE_WATCHER_SLEEP_INTERVAL = 1.0
+
+# For testing only - do not use in production
+_TEST_SLEEP_INTERVAL: float | None = None
+
+
 def watch_modules(
     graph: dataflow.DirectedGraph,
     reloader: ModuleReloader,
@@ -131,6 +137,7 @@ def watch_modules(
     # work with a copy to avoid race conditions
     # in CPython, dict.copy() is atomic
     sys_modules = sys.modules.copy()
+    sleep_interval = _TEST_SLEEP_INTERVAL or MODULE_WATCHER_SLEEP_INTERVAL
     while not should_exit.is_set():
         # Collect the modules used by each cell
         modules: dict[str, types.ModuleType] = {}
@@ -186,7 +193,7 @@ def watch_modules(
         # Don't proceed until enqueue_run_stale_cells() has been processed,
         # ie until stale cells have been rerun
         run_is_processed.wait()
-        time.sleep(1)
+        time.sleep(sleep_interval)
         # Update our snapshot of sys.modules
         sys_modules = sys.modules.copy()
 

tests/_runtime/reload/reload_test_utils.py

@@ -23,5 +23,5 @@ def update_file(path: pathlib.Path, code: str) -> None:
     (because that is stored in the file).  The only reliable way
     to achieve this seems to be to sleep.
     """
-    time.sleep(1.5)
+    time.sleep(1.05)
     path.write_text(textwrap.dedent(code))

tests/_runtime/reload/test_module_watcher.py

@@ -16,6 +16,19 @@
 from marimo._runtime.runtime import Kernel
 from tests.conftest import ExecReqProvider
 
+INTERVAL = 0.1
+
+
+@pytest.fixture(autouse=True)
+def _setup_test_sleep():
+    """Automatically set up faster sleep interval for all tests in this module"""
+    import marimo._runtime.reload.module_watcher as mw
+
+    old_interval = mw._TEST_SLEEP_INTERVAL
+    mw._TEST_SLEEP_INTERVAL = INTERVAL
+    yield
+    mw._TEST_SLEEP_INTERVAL = old_interval
+
 
 # these tests use random filenames for modules because they share
 # the same sys.modules object, and each test needs fresh modules
@@ -57,7 +70,7 @@ def foo():
     )
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(2.5)
+    await asyncio.sleep(INTERVAL * 3)
     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
@@ -117,7 +130,7 @@ def foo():
     )
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(2.5)
+    await asyncio.sleep(INTERVAL * 3)
     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
@@ -172,7 +185,7 @@ def foo():
     update_file(nested_module, "func = lambda : 2")
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(2.5)
+    await asyncio.sleep(INTERVAL * 3)
     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
@@ -227,7 +240,7 @@ def foo():
     update_file(nested_module, "func = lambda : 2")
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(2.5)
+    await asyncio.sleep(INTERVAL * 3)
     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
@@ -334,7 +347,7 @@ async def test_reload_package(
     update_file(nested_module, "func = lambda : 2")
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(2.5)
+    await asyncio.sleep(INTERVAL * 3)
     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
@@ -390,7 +403,7 @@ def foo():
     )
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(2.5)
+    await asyncio.sleep(INTERVAL * 3)
     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
@@ -439,7 +452,7 @@ def foo():
     )
 
     # wait for the watcher to pick up the change
-    await asyncio.sleep(2.5)
+    await asyncio.sleep(INTERVAL * 3)
     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
@@ -502,10 +515,10 @@ def foo():
     )
 
     # wait for the watcher to pick up the change
-    elapsed = 0
-    while elapsed < 10:
-        await asyncio.sleep(1)
-        elapsed += 1
+    retries = 0
+    while retries < 10:
+        await asyncio.sleep(INTERVAL)
+        retries += 1
         if k.graph.cells[er_1.cell_id].stale:
             break
 

tests/conftest.py

@@ -138,6 +138,8 @@ def teardown(self) -> None:
         teardown_context()
         self.stdout._watcher.stop()
         self.stderr._watcher.stop()
+        if self.k.module_watcher is not None:
+            self.k.module_watcher.stop()
         sys.modules["__main__"] = self._main