feat(api): JSON protocol pipette loading support in Protocol Engine

api/src/opentrons/protocol_api_experimental/protocol_context.py

@@ -47,6 +47,7 @@ def load_pipette(  # noqa: D102
         result = self._engine_client.load_pipette(
             pipette_name=PipetteName(pipette_name),
             mount=Mount(mount),
+            pipette_id=None,
         )
 
         return PipetteContext(

api/src/opentrons/protocol_engine/clients/sync_client.py

@@ -1,6 +1,6 @@
 """Synchronous ProtocolEngine client module."""
 from uuid import uuid4
-from typing import cast
+from typing import cast, Optional
 
 from opentrons.types import MountType
 
@@ -51,11 +51,13 @@ def load_pipette(
         self,
         pipette_name: PipetteName,
         mount: MountType,
+        pipette_id: Optional[str],
     ) -> commands.LoadPipetteResult:
         """Execute a LoadPipetteRequest and return the result."""
         request = commands.LoadPipetteRequest(
             pipetteName=pipette_name,
             mount=mount,
+            pipetteId=pipette_id
         )
         result = self._transport.execute_command(
             request=request,

api/src/opentrons/protocol_engine/commands/load_pipette.py

@@ -1,5 +1,8 @@
 """Load pipette command request, result, and implementation models."""
 from __future__ import annotations
+
+from typing import Optional
+
 from pydantic import BaseModel, Field
 
 from opentrons.types import MountType
@@ -19,6 +22,11 @@ class LoadPipetteRequest(BaseModel):
         ...,
         description="The mount the pipette should be present on.",
     )
+    pipetteId: Optional[str] = Field(
+        ...,
+        description="An optional ID to assign to this pipette. If None, an ID "
+                    "will be generated."
+    )
 
     def get_implementation(self) -> LoadPipetteImplementation:
         """Get the load pipette request's command implementation."""
@@ -44,6 +52,7 @@ async def execute(self, handlers: CommandHandlers) -> LoadPipetteResult:
         loaded_pipette = await handlers.equipment.load_pipette(
             pipette_name=self._request.pipetteName,
             mount=self._request.mount,
+            pipette_id=self._request.pipetteId
         )
 
         return LoadPipetteResult(pipetteId=loaded_pipette.pipette_id)

api/src/opentrons/protocol_engine/execution/equipment.py

@@ -1,6 +1,6 @@
 """Equipment command side-effect logic."""
 from dataclasses import dataclass
-from typing import Tuple
+from typing import Tuple, Optional
 
 from opentrons_shared_data.labware.dev_types import LabwareDefinition
 from opentrons.types import MountType
@@ -77,8 +77,19 @@ async def load_pipette(
         self,
         pipette_name: PipetteName,
         mount: MountType,
+        pipette_id: Optional[str],
     ) -> LoadedPipette:
-        """Ensure the requested pipette is attached."""
+        """Ensure the requested pipette is attached.
+
+        Args:
+            pipette_name: The pipette name.
+            mount: The mount on which pipette must be attached.
+            pipette_id: An optional identifier to assign the pipette. If None, an
+                identifier will be generated.
+
+        Returns:
+            A LoadedPipette object.
+        """
         other_mount = mount.other_mount()
         other_pipette = self._state.pipettes.get_pipette_data_by_mount(
             other_mount,
@@ -97,6 +108,7 @@ async def load_pipette(
         except RuntimeError as e:
             raise FailedToLoadPipetteError(str(e)) from e
 
-        pipette_id = self._resources.id_generator.generate_id()
+        pipette_id = pipette_id if pipette_id else \
+            self._resources.id_generator.generate_id()
 
         return LoadedPipette(pipette_id=pipette_id)

api/tests/opentrons/protocol_api_experimental/test_protocol_context.py

@@ -50,6 +50,7 @@ def test_load_pipette(
         engine_client.load_pipette(
             pipette_name=PipetteName.P300_SINGLE,
             mount=Mount.LEFT,
+            pipette_id=None,
         )
     ).then_return(commands.LoadPipetteResult(pipetteId="pipette-id"))
 
@@ -71,13 +72,15 @@ def test_load_instrument(
         engine_client.load_pipette(
             pipette_name=PipetteName.P300_MULTI,
             mount=Mount.LEFT,
+            pipette_id=None
         )
     ).then_return(commands.LoadPipetteResult(pipetteId="left-pipette-id"))
 
     decoy.when(
         engine_client.load_pipette(
             pipette_name=PipetteName.P300_SINGLE,
             mount=Mount.RIGHT,
+            pipette_id=None
         )
     ).then_return(commands.LoadPipetteResult(pipetteId="right-pipette-id"))
 

api/tests/opentrons/protocol_engine/clients/test_sync_client.py

@@ -92,6 +92,7 @@ def test_load_pipette(
     request = commands.LoadPipetteRequest(
         pipetteName=PipetteName.P300_SINGLE,
         mount=MountType.RIGHT,
+        pipetteId=None,
     )
 
     expected_result = commands.LoadPipetteResult(pipetteId="abc123")
@@ -103,6 +104,7 @@ def test_load_pipette(
     result = subject.load_pipette(
         pipette_name=PipetteName.P300_SINGLE,
         mount=MountType.RIGHT,
+        pipette_id=None,
     )
 
     assert result == expected_result

api/tests/opentrons/protocol_engine/commands/test_load_pipette.py

@@ -15,10 +15,12 @@ def test_load_pipette_request() -> None:
     request = LoadPipetteRequest(
         pipetteName=PipetteName.P300_SINGLE,
         mount=MountType.LEFT,
+        pipetteId=None
     )
 
     assert request.pipetteName == "p300_single"
     assert request.mount == MountType.LEFT
+    assert request.pipetteId is None
 
 
 def test_load_pipette_result() -> None:
@@ -37,6 +39,7 @@ async def test_load_pipette_implementation(mock_handlers: AsyncMock) -> None:
     request = LoadPipetteRequest(
         pipetteName=PipetteName.P300_SINGLE,
         mount=MountType.LEFT,
+        pipetteId="some id"
     )
     impl = request.get_implementation()
     result = await impl.execute(mock_handlers)
@@ -45,4 +48,5 @@ async def test_load_pipette_implementation(mock_handlers: AsyncMock) -> None:
     mock_handlers.equipment.load_pipette.assert_called_with(
         pipette_name="p300_single",
         mount=MountType.LEFT,
+        pipette_id="some id",
     )

api/tests/opentrons/protocol_engine/execution/test_equipment_handler.py

@@ -110,12 +110,29 @@ async def test_load_pipette_assigns_id(
     res = await handler.load_pipette(
         pipette_name=PipetteName.P300_SINGLE,
         mount=MountType.LEFT,
+        pipette_id=None,
     )
 
     assert type(res) == LoadedPipette
     assert res.pipette_id == "unique-id"
 
 
+async def test_load_pipette_uses_provided_id(
+    mock_resources_with_data: AsyncMock,
+    handler: EquipmentHandler,
+) -> None:
+    """It should use the provided ID rather than generating an ID for the pipette."""
+    res = await handler.load_pipette(
+        pipette_name=PipetteName.P300_SINGLE,
+        mount=MountType.LEFT,
+        pipette_id="my pipette id"
+    )
+
+    assert type(res) == LoadedPipette
+    assert res.pipette_id == "my pipette id"
+    mock_resources_with_data.id_generator.generate_id.assert_not_called()
+
+
 async def test_load_pipette_checks_checks_existence(
     mock_state_view: MagicMock,
     mock_hardware: AsyncMock,
@@ -126,6 +143,7 @@ async def test_load_pipette_checks_checks_existence(
     await handler.load_pipette(
         pipette_name=PipetteName.P300_SINGLE,
         mount=MountType.LEFT,
+        pipette_id=None,
     )
 
     mock_state_view.pipettes.get_pipette_data_by_mount.assert_called_with(
@@ -151,6 +169,7 @@ async def test_load_pipette_checks_checks_existence_with_already_loaded(
     await handler.load_pipette(
         pipette_name=PipetteName.P300_SINGLE,
         mount=MountType.RIGHT,
+        pipette_id=None,
     )
 
     mock_state_view.pipettes.get_pipette_data_by_mount.assert_called_with(
@@ -181,4 +200,5 @@ async def test_load_pipette_raises_if_pipette_not_attached(
         await handler.load_pipette(
             pipette_name=PipetteName.P300_SINGLE,
             mount=MountType.LEFT,
+            pipette_id=None,
         )

api/tests/opentrons/protocol_engine/state/test_pipette_state.py

@@ -25,6 +25,7 @@ def load_pipette_command(now: datetime) -> CompletedLoadLabware:
         request=cmd.LoadPipetteRequest(
             pipetteName=PipetteName.P300_SINGLE,
             mount=MountType.LEFT,
+            pipetteId=None,
         ),
         result=cmd.LoadPipetteResult(pipetteId="pipette-id"),
         created_at=now,

robot-server/tests/integration/sessions/test_live_protocol.tavern.yaml

@@ -86,6 +86,7 @@ stages:
           data: &create_pipette_data
             pipetteName: p300_single
             mount: right
+            pipetteId: null
     response:
       status_code: 200
       json:
@@ -100,6 +101,31 @@ stages:
           completedAt: !anystr
           result:
             pipetteId: !anystr
+  - name: Load pipette command specifying the id
+    request:
+      url: "{host:s}:{port:d}/sessions/{session_id}/commands/execute"
+      method: POST
+      json:
+        data:
+          command: equipment.loadPipette
+          data: &create_pipette_data_with_id
+            pipetteName: p10_single
+            mount: left
+            pipetteId: my pipette
+    response:
+      status_code: 200
+      json:
+        links: !anydict
+        data:
+          id: !anystr
+          data: *create_pipette_data_with_id
+          command: equipment.loadPipette
+          status: executed
+          createdAt: !anystr
+          startedAt: !anystr
+          completedAt: !anystr
+          result:
+            pipetteId: my pipette
   - name: Delete the session
     request:
       url: "{host:s}:{port:d}/sessions/{session_id}"

robot-server/tests/service/session/models/test_command.py

@@ -53,15 +53,17 @@ def test_not_empty():
             "command": "equipment.loadPipette",
             "data": {
                 "pipetteName": "p10_single",
-                "mount": "left"
+                "mount": "left",
+                "pipetteId": None,
             }
         }
     })
     assert request.data.command == \
            command_definitions.EquipmentCommand.load_pipette
     assert request.data.data == pe_commands.LoadPipetteRequest(
         pipetteName="p10_single",
-        mount=MountType.LEFT
+        mount=MountType.LEFT,
+        pipetteId=None,
     )
 
     dt = datetime(2000, 1, 1)
@@ -81,7 +83,8 @@ def test_not_empty():
            command_definitions.EquipmentCommand.load_pipette
     assert response.data == pe_commands.LoadPipetteRequest(
         pipetteName="p10_single",
-        mount=MountType.LEFT
+        mount=MountType.LEFT,
+        pipetteId=None,
     )
     assert response.id == "id"
     assert response.createdAt == dt

robot-server/tests/service/session/session_types/live_protocol/test_command_executor.py

@@ -136,6 +136,7 @@ async def test_load_instrument(command_executor, mock_protocol_engine):
     request_body = pe_commands.LoadPipetteRequest(
         pipetteName="p10_single",
         mount=MountType.LEFT,
+        pipetteId=None,
     )
 
     protocol_engine_response = pe_commands.CompletedCommand(