Prevent race condition by enabling page events before navigation

src/py/kaleido/_kaleido_tab/_tab.py

@@ -78,16 +78,16 @@ async def navigate(self, url: str | Path = ""):
         # Apply headers if they exist
         await self._apply_headers()
 
+        # Enable Page events before navigation; completed load events are not replayed.
+        _logger.debug2(f"Calling Page.enable on {self.tab}")
+        _raise_error(await self.tab.send_command("Page.enable"))
+
         # Navigating page. This will trigger the above events.
         _logger.debug2(f"Calling Page.navigate on {self.tab}")
         _raise_error(await self.tab.send_command("Page.navigate", params={"url": url}))
 
-        # Enabling page events (for page_ready- like all events, if already
-        # ready, the latest will fire immediately)
-        _logger.debug2(f"Calling Page.enable on {self.tab}")
-        _raise_error(await self.tab.send_command("Page.enable"))
-
-        # Enabling javascript events (for javascript_ready)
+        # Enabling javascript events (for javascript_ready).
+        # We do this after navigation to receive the new execution context.
         _logger.debug2(f"Calling Runtime.enable on {self.tab}")
         _raise_error(await self.tab.send_command("Runtime.enable"))
 

src/py/tests/test_kaleido_tab.py

@@ -0,0 +1,39 @@
+import asyncio
+from unittest.mock import AsyncMock, Mock
+
+import choreographer
+
+from kaleido._kaleido_tab import _KaleidoTab
+
+
+async def test_navigate_when_page_loads_before_navigation_returns():
+    tab = Mock(spec=choreographer.Tab)
+    javascript_ready = asyncio.get_running_loop().create_future()
+    page_ready = asyncio.get_running_loop().create_future()
+    tab.subscribe_once.side_effect = {
+        "Runtime.executionContextCreated": javascript_ready,
+        "Page.loadEventFired": page_ready,
+    }.__getitem__
+    page_enabled = False
+    navigated_context_id = 2
+    context_id = 1
+
+    async def send_command(command, params=None):  # noqa: ARG001
+        nonlocal page_enabled, context_id
+        if command == "Page.enable":
+            page_enabled = True
+        elif command == "Page.navigate":
+            context_id = navigated_context_id
+            # Chrome does not replay a load event missed while Page was disabled.
+            if page_enabled:
+                page_ready.set_result({})
+        elif command == "Runtime.enable":
+            javascript_ready.set_result({"params": {"context": {"id": context_id}}})
+        return {"result": {}}
+
+    tab.send_command = AsyncMock(side_effect=send_command)
+    kaleido_tab = _KaleidoTab(tab)
+
+    await asyncio.wait_for(kaleido_tab.navigate("file:///index.html"), timeout=1)
+
+    assert kaleido_tab._current_js_id == navigated_context_id  # noqa: SLF001