Centralize the logic to pull until the run completes.

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

@@ -1,4 +1,3 @@
-import asyncio
 from dataclasses import dataclass, field
 from pathlib import Path
 from shutil import copytree
@@ -9,14 +8,13 @@
 import pytest
 
 from tests.integration.dev_server import DevServer
-from tests.integration.robot_client import RobotClient
+from tests.integration.robot_client import RobotClient, poll_until_run_completes
 
 from .persistence_snapshots_dir import PERSISTENCE_SNAPSHOTS_DIR
 
 # Allow plenty of time for database migrations, which can take a while in our CI runners.
 _STARTUP_TIMEOUT = 60
 
-_POLL_INTERVAL = 0.1
 _RUN_TIMEOUT = 5
 
 # Our Tavern tests have servers that stay up for the duration of the test session.
@@ -193,15 +191,9 @@ async def test_rerun_flex_dev_compat() -> None:
  )
 
  with anyio.fail_after(_RUN_TIMEOUT):
- final_status = await _poll_until_not_running(client, new_run["id"])
+ final_status = (
+ await poll_until_run_completes(
+ robot_client=client, run_id=new_run["id"]
+ )
+ )["data"]["status"]
  assert final_status == "succeeded"
-
-
-async def _poll_until_not_running(robot_client: RobotClient, run_id: str) -> str:
- while True:
- latest_status = (await robot_client.get_run(run_id)).json()["data"]["status"]
- if latest_status not in {"running", "finishing"}:
- return latest_status # type: ignore[no-any-return]
- else:
- # Sleep, then poll again.
- await asyncio.sleep(_POLL_INTERVAL)

robot-server/tests/integration/http_api/runs/test_labware_offsets_on_compatible_modules.py

@@ -1,11 +1,10 @@
-import asyncio
 from textwrap import dedent
 from typing import Any, AsyncGenerator
 
 import anyio
 import pytest
 
-from tests.integration.robot_client import RobotClient
+from tests.integration.robot_client import RobotClient, poll_until_run_completes
 
 
 # An arbitrary choice of labware.
@@ -37,16 +36,11 @@ async def poll_until_run_succeeds(robot_client: RobotClient, run_id: str) -> Any
 
  Return the completed run response.
  """
- completed_run_statuses = {"stopped", "failed", "succeeded"}
- while True:
- run = (await robot_client.get_run(run_id=run_id)).json()
- status = run["data"]["status"]
- if status in completed_run_statuses:
- assert status == "succeeded"
- return run
- else:
- # The run is still ongoing. Wait a beat, then poll again.
- await asyncio.sleep(RUN_POLL_INTERVAL)
+ with anyio.fail_after(RUN_POLL_TIMEOUT):
+ final_status = (
+ await poll_until_run_completes(robot_client=robot_client, run_id=run_id)
+ )["data"]["status"]
+ assert final_status == "succeeded"
 
 
 @pytest.fixture
@@ -144,8 +138,7 @@ async def test_labware_offsets_on_compatible_modules(
  await robot_client.post_run_action(
  run_id=run_id, req_body={"data": {"actionType": "play"}}
  )
- with anyio.fail_after(RUN_POLL_TIMEOUT):
- await poll_until_run_succeeds(robot_client=robot_client, run_id=run_id)
+ await poll_until_run_succeeds(robot_client=robot_client, run_id=run_id)
 
  # Retrieve details about the protocol's completed commands.
  commands = (await robot_client.get_run_commands(run_id=run_id)).json()

robot-server/tests/integration/robot_client.py

@@ -13,6 +13,8 @@
 STARTUP_WAIT = 20
 SHUTDOWN_WAIT = 20
 
+_RUN_POLL_INTERVAL = 0.1
+
 
 class RobotClient:
  """Client for the robot's HTTP API.
@@ -132,9 +134,7 @@ async def post_protocol(
  multipart_upload_name = "files"
 
  with contextlib.ExitStack() as file_exit_stack:
- opened_files: List[
- Union[BinaryIO, Tuple[str, bytes]],
- ] = []
+ opened_files: List[Union[BinaryIO, Tuple[str, bytes]],] = []
 
  for file in files:
  if isinstance(file, Path):
@@ -317,3 +317,26 @@ async def delete_session(self, session_id: str) -> Response:
  )
  response.raise_for_status()
  return response
+
+
+async def poll_until_run_completes(
+ robot_client: RobotClient, run_id: str, poll_interval: float = _RUN_POLL_INTERVAL
+) -> Any:
+ """Wait until a run completes.
+
+ You probably want to wrap this in an `anyio.fail_after()` timeout in case something causes
+ the run to hang forever.
+
+ Returns:
+ The completed run response. You can inspect its `status` to see whether it
+ succeeded, failed, or was stopped.
+ """
+ completed_run_statuses = {"stopped", "failed", "succeeded"}
+ while True:
+ run = (await robot_client.get_run(run_id=run_id)).json()
+ status = run["data"]["status"]
+ if status in completed_run_statuses:
+ return run
+ else:
+ # The run is still ongoing. Wait a beat, then poll again.
+ await asyncio.sleep(poll_interval)