await add and wait for command

Opentrons · May 30, 2024 · 8e06362 · 8e06362
1 parent 9531aa7
commit 8e06362
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/api/src/opentrons/protocol_runner/run_orchestrator.py b/api/src/opentrons/protocol_runner/run_orchestrator.py
@@ -102,7 +102,7 @@ def build_orchestrator(
             protocol_engine=protocol_engine,
         )
 
-    def add_command_and_wait_for_interval(self, command: CommandCreate, wait_until_complete: bool = False,
+    async def add_command_and_wait_for_interval(self, command: CommandCreate, wait_until_complete: bool = False,
                                           timeout: Optional[int] = None, failed_command_id: Optional[str] = None) -> Command:
         added_command = self._protocol_engine.add_command(request=command, failed_command_id=failed_command_id)
         if wait_until_complete:

diff --git a/robot-server/robot_server/runs/engine_store.py b/robot-server/robot_server/runs/engine_store.py
@@ -395,14 +395,14 @@ def add_labware_offset(self, request: LabwareOffsetCreate) -> LabwareOffset:
     def add_labware_definition(self, definition: LabwareDefinition) -> LabwareUri:
         return self._run_orchestrator.engine.add_labware_definition(definition)
 
-    def add_command_and_wait_for_interval(
+    async def add_command_and_wait_for_interval(
         self,
         request: CommandCreate,
         wait_until_complete: bool = False,
         timeout: Optional[int] = None,
         failed_command_id: Optional[str] = None,
     ) -> Command:
-        return self._run_orchestrator.add_command_and_wait_for_interval(
+        return await self._run_orchestrator.add_command_and_wait_for_interval(
             command=request,
             failed_command_id=failed_command_id,
             wait_until_complete=wait_until_complete,

diff --git a/robot-server/robot_server/runs/router/commands_router.py b/robot-server/robot_server/runs/router/commands_router.py
@@ -208,7 +208,7 @@ async def create_run_command(
     command_intent = request_body.data.intent or pe_commands.CommandIntent.SETUP
     command_create = request_body.data.copy(update={"intent": command_intent})
     try:
-        command = engine_store.add_command_and_wait_for_interval(
+        command = await engine_store.add_command_and_wait_for_interval(
             request=command_create,
             failed_command_id=failedCommandId,
             wait_until_complete=waitUntilComplete,

diff --git a/robot-server/tests/runs/router/test_commands_router.py b/robot-server/tests/runs/router/test_commands_router.py
@@ -114,7 +114,7 @@ def _stub_queued_command_state(*_a: object, **_k: object) -> pe_commands.Command
         return command_once_added
 
     decoy.when(
-        mock_engine_store.add_command_and_wait_for_interval(
+        await mock_engine_store.add_command_and_wait_for_interval(
             request=pe_commands.WaitForResumeCreate(
                 params=pe_commands.WaitForResumeParams(message="Hello"),
                 intent=pe_commands.CommandIntent.SETUP,
@@ -201,7 +201,7 @@ def _stub_completed_command_state(*_a: object, **_k: object) -> None:
         )
 
     decoy.when(
-        mock_engine_store.add_command_and_wait_for_interval(
+        await mock_engine_store.add_command_and_wait_for_interval(
             request=command_request, failed_command_id=None
         )
     ).then_do(_stub_queued_command_state)