implement labware handling

api/src/opentrons/protocol_engine/state/labware.py

@@ -31,6 +31,7 @@
     Command,
     LoadLabwareResult,
     MoveLabwareResult,
+    ReloadLabwareResult,
 )
 from ..types import (
     DeckSlotLocation,
@@ -175,7 +176,7 @@ def handle_action(self, action: Action) -> None:
 
     def _handle_command(self, command: Command) -> None:
         """Modify state in reaction to a command."""
-        if isinstance(command.result, LoadLabwareResult):
+        if isinstance(command.result, (LoadLabwareResult, ReloadLabwareResult)):
             # If the labware load refers to an offset, that offset must actually exist.
             if command.result.offsetId is not None:
                 assert command.result.offsetId in self._state.labware_offsets_by_id
@@ -187,12 +188,16 @@ def _handle_command(self, command: Command) -> None:
             )
 
             self._state.definitions_by_uri[definition_uri] = command.result.definition
+            if isinstance(command.result, LoadLabwareResult):
+                location = command.params.location
+            else:
+                location = self._state.labware_by_id[command.result.labwareId].location
 
             self._state.labware_by_id[
                 command.result.labwareId
             ] = LoadedLabware.construct(
                 id=command.result.labwareId,
-                location=command.params.location,
+                location=location,
                 loadName=command.result.definition.parameters.loadName,
                 definitionUri=definition_uri,
                 offsetId=command.result.offsetId,

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

@@ -577,3 +577,34 @@ def create_prepare_to_aspirate_command(pipette_id: str) -> cmd.PrepareToAspirate
         params=params,
         result=result,
     )
+
+
+def create_reload_labware_command(
+    labware_id: str,
+    definition: LabwareDefinition,
+    offset_id: Optional[str],
+    display_name: Optional[str],
+) -> cmd.ReloadLabware:
+    """Create a completed ReloadLabware command."""
+    params = cmd.ReloadLabwareParams(
+        loadName=definition.parameters.loadName,
+        namespace=definition.namespace,
+        version=definition.version,
+        labwareId=labware_id,
+        displayName=display_name,
+    )
+
+    result = cmd.ReloadLabwareResult(
+        labwareId=labware_id,
+        definition=definition,
+        offsetId=offset_id,
+    )
+
+    return cmd.ReloadLabware(
+        id="command-id",
+        key="command-key",
+        status=cmd.CommandStatus.SUCCEEDED,
+        createdAt=datetime.now(),
+        params=params,
+        result=result,
+    )

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

@@ -28,6 +28,7 @@
 from .command_fixtures import (
     create_load_labware_command,
     create_move_labware_command,
+    create_reload_labware_command,
 )
 
 
@@ -132,6 +133,73 @@ def test_handles_load_labware(
     assert subject.state.definitions_by_uri[expected_definition_uri] == well_plate_def
 
 
+def test_handles_reload_labware(
+    subject: LabwareStore,
+    well_plate_def: LabwareDefinition,
+    flex_50uL_tiprack: LabwareDefinition,
+) -> None:
+    """It should override labware data in the state."""
+    load_labware = create_load_labware_command(
+        location=DeckSlotLocation(slotName=DeckSlotName.SLOT_A1),
+        labware_id="test-labware-id",
+        definition=well_plate_def,
+        display_name="display-name",
+        offset_id=None,
+    )
+
+    subject.handle_action(
+        SucceedCommandAction(private_result=None, command=load_labware)
+    )
+    assert subject.state.labware_by_id[
+        "test-labware-id"
+    ].definitionUri == uri_from_details(
+        load_name=well_plate_def.parameters.loadName,
+        namespace=well_plate_def.namespace,
+        version=well_plate_def.version,
+    )
+
+    offset_request = LabwareOffsetCreate(
+        definitionUri="offset-definition-uri",
+        location=LabwareOffsetLocation(slotName=DeckSlotName.SLOT_1),
+        vector=LabwareOffsetVector(x=1, y=2, z=3),
+    )
+    subject.handle_action(
+        AddLabwareOffsetAction(
+            request=offset_request,
+            labware_offset_id="offset-id",
+            created_at=datetime(year=2021, month=1, day=2),
+        )
+    )
+    reload_labware = create_reload_labware_command(
+        labware_id="test-labware-id",
+        definition=flex_50uL_tiprack,
+        offset_id="offset-id",
+        display_name="display-name-2",
+    )
+    subject.handle_action(
+        SucceedCommandAction(private_result=None, command=reload_labware)
+    )
+
+    expected_definition_uri = uri_from_details(
+        load_name=flex_50uL_tiprack.parameters.loadName,
+        namespace=flex_50uL_tiprack.namespace,
+        version=flex_50uL_tiprack.version,
+    )
+
+    expected_labware_data = LoadedLabware(
+        id="test-labware-id",
+        loadName=flex_50uL_tiprack.parameters.loadName,
+        definitionUri=expected_definition_uri,
+        location=DeckSlotLocation(slotName=DeckSlotName.SLOT_A1),
+        offsetId="offset-id",
+        displayName="display-name-2",
+    )
+    assert subject.state.labware_by_id["test-labware-id"] == expected_labware_data
+    assert (
+        subject.state.definitions_by_uri[expected_definition_uri] == flex_50uL_tiprack
+    )
+
+
 def test_handles_add_labware_definition(
     subject: LabwareStore,
     well_plate_def: LabwareDefinition,