From 0159b439cb508ee2b2e3583d1c9fa7bff7a7f7d5 Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Fri, 17 Oct 2025 11:47:11 +0200
Subject: [PATCH] Make test_execute more robust

---
 pyproject.toml        |  1 +
 tests/test_execute.py |  7 ++++++-
 tests/test_server.py  | 12 ++++++++++--
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ae114e7d..643ca1f1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -65,6 +65,7 @@ test = [
     "types-setuptools",
     "pytest",
     "pytest-timeout",
+    "pytest-rerunfailures",
     "pytest-env",
     "httpx",
     "httpx-ws >=0.4.1",
diff --git a/tests/test_execute.py b/tests/test_execute.py
index 850633be..38934f2f 100644
--- a/tests/test_execute.py
+++ b/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(
diff --git a/tests/test_server.py b/tests/test_server.py
index 7001ed89..4cab44db 100644
--- a/tests/test_server.py
+++ b/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}'