Make test_execute more robust

Author: davidbrochart

pyproject.toml

@@ -65,6 +65,7 @@ test = [
     "types-setuptools",
     "pytest",
     "pytest-timeout",
+    "pytest-rerunfailures",
     "pytest-env",
     "httpx",
     "httpx-ws >=0.4.1",

tests/test_execute.py

@@ -86,6 +86,7 @@ async def recv(self) -> bytes:
 
 
 @pytest.mark.anyio
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("auth_mode", ("noauth",))
 async def test_execute(auth_mode, unused_tcp_port):
     url = f"http://127.0.0.1:{unused_tcp_port}"
@@ -152,7 +153,11 @@ def callback(aevent, events, event):
         ):
             # connect to the shared notebook document
             # wait for file to be loaded and Y model to be created in server and client
-            await anyio.sleep(0.5)
+            with anyio.fail_after(1):
+                while True:
+                    if len(ynb.ycells) >= 2:
+                        break
+                    await anyio.sleep(0.1)
             # execute notebook
             for cell_idx in range(2):
                 response = await http.post(

tests/test_server.py

@@ -175,7 +175,11 @@ def callback(aevent, events, event):
     ):
         # connect to the shared notebook document
         # wait for file to be loaded and Y model to be created in server and client
-        await anyio.sleep(0.5)
+        with anyio.fail_after(1):
+            while True:
+                if len(ynb.ycells) >= 2:
+                    break
+                await anyio.sleep(0.1)
         # execute notebook
         for cell_idx in range(2):
             response = requests.post(