Make test_execute more robust

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:
+                    await anyio.sleep(0.1)
+                    if len(ynb.ycells) >= 2:
+                        break
             # 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:
+                await anyio.sleep(0.1)
+                if len(ynb.ycells) >= 2:
+                    break
         # execute notebook
         for cell_idx in range(2):
             response = requests.post(
@@ -215,10 +219,14 @@ async def connect_ywidget(url, guid):
         aconnect_ws(f"{url}/api/collaboration/room/ywidget:{guid}") as websocket,
         WebsocketProvider(ywidget_doc, Websocket(websocket, guid)),
     ):
-        await anyio.sleep(0.5)
         attrs = Map()
         model_name = Text()
         ywidget_doc["_attrs"] = attrs
         ywidget_doc["_model_name"] = model_name
+        with anyio.fail_after(1):
+            while True:
+                await anyio.sleep(0.1)
+                if len(attrs) > 0:
+                    break
         assert str(model_name) == "Switch"
         assert str(attrs) == '{"value":true}'