chore(robot-server): Work around tests hanging because server shutdown is hanging

robot-server/tests/integration/http_api/persistence/test_reset.py

@@ -4,10 +4,11 @@
 from shutil import copytree
 from tempfile import TemporaryDirectory
 
+import anyio
 import httpx
 
 from tests.integration.dev_server import DevServer
-from tests.integration.robot_client import RobotClient
+from tests.integration.robot_client import RobotClient, poll_until_all_analyses_complete
 from tests.integration.protocol_files import get_py_protocol, get_json_protocol
 
 from .persistence_snapshots_dir import PERSISTENCE_SNAPSHOTS_DIR
@@ -98,6 +99,13 @@ async def test_upload_protocols_and_reset_persistence_dir() -> None:
             with get_json_protocol(secrets.token_urlsafe(16)) as file:
                 await robot_client.post_protocol([Path(file.name)])
 
+            with anyio.fail_after(30):
+                # todo(mm, 2024-09-20): This works around a bug where robot-server
+                # shutdown will hang if there is an ongoing analysis. This slows down
+                # this test and should be removed when that bug is fixed.
+                # https://opentrons.atlassian.net/browse/EXEC-716
+                await poll_until_all_analyses_complete(robot_client)
+
             await robot_client.post_setting_reset_options({"runsHistory": True})
 
             result = await robot_client.get_protocols()

robot-server/tests/integration/http_api/protocols/test_persistence.py

@@ -3,12 +3,13 @@
 import secrets
 from typing import Callable, Dict, IO, List
 
+import anyio
 import pytest
 
 from robot_server.persistence.file_and_directory_names import LATEST_VERSION_DIRECTORY
 
 from tests.integration.dev_server import DevServer
-from tests.integration.robot_client import RobotClient
+from tests.integration.robot_client import RobotClient, poll_until_all_analyses_complete
 from tests.integration.protocol_files import get_py_protocol, get_json_protocol
 
 
@@ -39,7 +40,7 @@ async def test_protocols_and_analyses_persist(
                     await robot_client.post_protocol([Path(file.name)])
 
             await asyncio.wait_for(
-                _wait_for_all_analyses_to_complete(robot_client), timeout=30
+                poll_until_all_analyses_complete(robot_client), timeout=30
             )
 
             # The protocols response will include analysis statuses. Fetch it
@@ -108,6 +109,13 @@ async def test_protocol_labware_files_persist() -> None:
             # we can ignore the whole field in this test to avoid the nondeterminism.
             del protocol_detail["analysisSummaries"]
 
+            with anyio.fail_after(30):
+                # todo(mm, 2024-09-20): This works around a bug where robot-server
+                # shutdown will hang if there is an ongoing analysis. This slows down
+                # this test and should be removed when that bug is fixed.
+                # https://opentrons.atlassian.net/browse/EXEC-716
+                await poll_until_all_analyses_complete(robot_client)
+
             server.stop()
             assert await robot_client.dead(), "Dev Robot did not stop."
             server.start()
@@ -168,16 +176,3 @@ async def _get_all_analyses(robot_client: RobotClient) -> Dict[str, List[object]
         analyses_by_protocol_id[protocol_id] = analyses_on_this_protocol
 
     return analyses_by_protocol_id
-
-
-async def _wait_for_all_analyses_to_complete(robot_client: RobotClient) -> None:
-    async def _all_analyses_are_complete() -> bool:
-        protocols = (await robot_client.get_protocols()).json()
-        for protocol in protocols["data"]:
-            for analysis_summary in protocol["analysisSummaries"]:
-                if analysis_summary["status"] != "completed":
-                    return False
-        return True
-
-    while not await _all_analyses_are_complete():
-        await asyncio.sleep(0.1)

robot-server/tests/integration/robot_client.py

@@ -14,6 +14,7 @@
 _SHUTDOWN_WAIT = 20
 
 _RUN_POLL_INTERVAL = 0.1
+_ANALYSIS_POLL_INTERVAL = 0.1
 
 
 class RobotClient:
@@ -398,3 +399,24 @@ async def poll_until_run_completes(
         else:
             # The run is still ongoing. Wait a beat, then poll again.
             await asyncio.sleep(poll_interval)
+
+
+async def poll_until_all_analyses_complete(
+    robot_client: RobotClient, poll_interval: float = _ANALYSIS_POLL_INTERVAL
+) -> None:
+    """Wait until all pending analyses have completed.
+
+    You probably want to wrap this in an `anyio.fail_after()` timeout in case something causes
+    an analysis to hang forever.
+    """
+
+    async def _all_analyses_are_complete() -> bool:
+        protocols = (await robot_client.get_protocols()).json()
+        for protocol in protocols["data"]:
+            for analysis_summary in protocol["analysisSummaries"]:
+                if analysis_summary["status"] != "completed":
+                    return False
+        return True
+
+    while not await _all_analyses_are_complete():
+        await asyncio.sleep(poll_interval)