From c7265da62a940e61ab128eba15be4097f4b9365e Mon Sep 17 00:00:00 2001
From: Laura Cox <31892318+Laura-Danielle@users.noreply.github.com>
Date: Thu, 22 Feb 2024 23:24:03 +0200
Subject: [PATCH] refactor(api): Save full nozzle map configuration and update
 state store accordingly (#14529)

To make a few operations easier in protocol engine, we will keep the
nozzle map in state always. I will think about this further in RSS-443,
but for now it should unblock other partial tip work.
---
 .../commands/configuring_common.py            |  5 ++-
 .../protocol_engine/execution/equipment.py    |  4 +--
 .../resources/pipette_data_provider.py        |  1 -
 .../protocol_engine/state/pipettes.py         | 15 ++++++--
 .../commands/test_configure_for_volume.py     |  3 +-
 .../commands/test_configure_nozzle_layout.py  |  9 ++---
 .../commands/test_load_pipette.py             |  3 +-
 .../execution/test_equipment_handler.py       |  3 +-
 .../protocol_engine/pipette_fixtures.py       | 34 +++++++++++++++++++
 .../resources/test_pipette_data_provider.py   |  6 ++--
 .../state/test_geometry_view.py               |  7 ++++
 .../state/test_pipette_store.py               |  5 +--
 .../state/test_pipette_view.py                |  1 +
 .../protocol_engine/state/test_tip_state.py   |  4 +--
 14 files changed, 72 insertions(+), 28 deletions(-)

diff --git a/api/src/opentrons/protocol_engine/commands/configuring_common.py b/api/src/opentrons/protocol_engine/commands/configuring_common.py
index ec5917d9931..17ffc2adef4 100644
--- a/api/src/opentrons/protocol_engine/commands/configuring_common.py
+++ b/api/src/opentrons/protocol_engine/commands/configuring_common.py
@@ -1,7 +1,6 @@
 """Common configuration command base models."""
 
 from pydantic import BaseModel, Field
-from typing import Optional
 from dataclasses import dataclass
 from opentrons.hardware_control.nozzle_manager import (
     NozzleMap,
@@ -22,7 +21,7 @@ class PipetteNozzleLayoutResultMixin(BaseModel):
     """A nozzle layout result for updating the pipette state."""
 
     pipette_id: str
-    nozzle_map: Optional[NozzleMap] = Field(
-        default=None,
+    nozzle_map: NozzleMap = Field(
+        ...,
         description="A dataclass object holding information about the current nozzle configuration.",
     )
diff --git a/api/src/opentrons/protocol_engine/execution/equipment.py b/api/src/opentrons/protocol_engine/execution/equipment.py
index c1ac272a64d..2487ad50aaa 100644
--- a/api/src/opentrons/protocol_engine/execution/equipment.py
+++ b/api/src/opentrons/protocol_engine/execution/equipment.py
@@ -202,7 +202,6 @@ async def load_pipette(
         )
 
         pipette_id = pipette_id or self._model_utils.generate_id()
-
         if not use_virtual_pipettes:
             cache_request = {mount.to_hw_mount(): pipette_name_value}
 
@@ -244,7 +243,6 @@ async def load_pipette(
                 )
             )
         serial = serial_number or ""
-
         return LoadedPipetteData(
             pipette_id=pipette_id,
             serial_number=serial,
@@ -390,7 +388,7 @@ async def configure_nozzle_layout(
         primary_nozzle: Optional[str] = None,
         front_right_nozzle: Optional[str] = None,
         back_left_nozzle: Optional[str] = None,
-    ) -> Optional[NozzleMap]:
+    ) -> NozzleMap:
         """Ensure the requested nozzle layout is compatible with the current pipette.
 
         Args:
diff --git a/api/src/opentrons/protocol_engine/resources/pipette_data_provider.py b/api/src/opentrons/protocol_engine/resources/pipette_data_provider.py
index 09b715c12e8..6264c2c526b 100644
--- a/api/src/opentrons/protocol_engine/resources/pipette_data_provider.py
+++ b/api/src/opentrons/protocol_engine/resources/pipette_data_provider.py
@@ -18,7 +18,6 @@
 )
 
 from ..types import FlowRates
-from ...types import Point
 
 
 @dataclass(frozen=True)
diff --git a/api/src/opentrons/protocol_engine/state/pipettes.py b/api/src/opentrons/protocol_engine/state/pipettes.py
index c326d7af6bd..26c71a6884f 100644
--- a/api/src/opentrons/protocol_engine/state/pipettes.py
+++ b/api/src/opentrons/protocol_engine/state/pipettes.py
@@ -111,6 +111,7 @@ class StaticPipetteConfig:
     nozzle_offset_z: float
     pipette_bounding_box_offsets: PipetteBoundingBoxOffsets
     bounding_nozzle_offsets: BoundingNozzlesOffsets
+    default_nozzle_map: NozzleMap
 
 
 @dataclass
@@ -180,11 +181,15 @@ def _handle_command(  # noqa: C901
                     front_right_corner=config.front_right_corner_offset,
                 ),
                 bounding_nozzle_offsets=BoundingNozzlesOffsets(
-                    back_left_offset=config.back_left_nozzle_offset,
-                    front_right_offset=config.front_right_nozzle_offset,
+                    back_left_offset=config.nozzle_map.back_left_nozzle_offset,
+                    front_right_offset=config.nozzle_map.front_right_nozzle_offset,
                 ),
+                default_nozzle_map=config.nozzle_map,
             )
             self._state.flow_rates_by_id[private_result.pipette_id] = config.flow_rates
+            self._state.nozzle_configuration_by_id[
+                private_result.pipette_id
+            ] = config.nozzle_map
         elif isinstance(private_result, PipetteNozzleLayoutResultMixin):
             self._state.nozzle_configuration_by_id[
                 private_result.pipette_id
@@ -201,7 +206,11 @@ def _handle_command(  # noqa: C901
             self._state.aspirated_volume_by_id[pipette_id] = None
             self._state.movement_speed_by_id[pipette_id] = None
             self._state.attached_tip_by_id[pipette_id] = None
-            self._state.nozzle_configuration_by_id[pipette_id] = None
+            static_config = self._state.static_config_by_id.get(pipette_id)
+            if static_config:
+                self._state.nozzle_configuration_by_id[
+                    pipette_id
+                ] = static_config.default_nozzle_map
 
         elif isinstance(command.result, (AspirateResult, AspirateInPlaceResult)):
             pipette_id = command.params.pipetteId
diff --git a/api/tests/opentrons/protocol_engine/commands/test_configure_for_volume.py b/api/tests/opentrons/protocol_engine/commands/test_configure_for_volume.py
index ebd57bd24c5..897b8330205 100644
--- a/api/tests/opentrons/protocol_engine/commands/test_configure_for_volume.py
+++ b/api/tests/opentrons/protocol_engine/commands/test_configure_for_volume.py
@@ -16,7 +16,8 @@
     ConfigureForVolumePrivateResult,
     ConfigureForVolumeImplementation,
 )
-from opentrons.types import Point
+from opentrons_shared_data.pipette.dev_types import PipetteNameType
+from ..pipette_fixtures import get_default_nozzle_map
 
 
 async def test_configure_for_volume_implementation(
diff --git a/api/tests/opentrons/protocol_engine/commands/test_configure_nozzle_layout.py b/api/tests/opentrons/protocol_engine/commands/test_configure_nozzle_layout.py
index f6b45376f7e..23cdddd98be 100644
--- a/api/tests/opentrons/protocol_engine/commands/test_configure_nozzle_layout.py
+++ b/api/tests/opentrons/protocol_engine/commands/test_configure_nozzle_layout.py
@@ -1,7 +1,7 @@
 """Test configure nozzle layout commands."""
 import pytest
 from decoy import Decoy
-from typing import Union, Optional, Dict
+from typing import Union, Dict
 from collections import OrderedDict
 
 from opentrons.protocol_engine.execution import (
@@ -73,11 +73,6 @@
             ),
             {"primary_nozzle": "A1", "front_right_nozzle": "E1"},
         ],
-        [
-            AllNozzleLayoutConfiguration(),
-            None,
-            {},
-        ],
     ],
 )
 async def test_configure_nozzle_layout_implementation(
@@ -90,7 +85,7 @@ async def test_configure_nozzle_layout_implementation(
         QuadrantNozzleLayoutConfiguration,
         SingleNozzleLayoutConfiguration,
     ],
-    expected_nozzlemap: Optional[NozzleMap],
+    expected_nozzlemap: NozzleMap,
     nozzle_params: Dict[str, str],
 ) -> None:
     """A ConfigureForVolume command should have an execution implementation."""
diff --git a/api/tests/opentrons/protocol_engine/commands/test_load_pipette.py b/api/tests/opentrons/protocol_engine/commands/test_load_pipette.py
index 428e2ad0672..c94391dac68 100644
--- a/api/tests/opentrons/protocol_engine/commands/test_load_pipette.py
+++ b/api/tests/opentrons/protocol_engine/commands/test_load_pipette.py
@@ -4,7 +4,7 @@
 
 from opentrons_shared_data.pipette.dev_types import PipetteNameType
 from opentrons_shared_data.robot.dev_types import RobotType
-from opentrons.types import MountType, Point
+from opentrons.types import MountType
 
 from opentrons.protocol_engine.errors import InvalidSpecificationForRobotTypeError
 from opentrons.protocol_engine.types import FlowRates
@@ -19,6 +19,7 @@
     LoadPipettePrivateResult,
     LoadPipetteImplementation,
 )
+from ..pipette_fixtures import get_default_nozzle_map
 
 
 async def test_load_pipette_implementation(
diff --git a/api/tests/opentrons/protocol_engine/execution/test_equipment_handler.py b/api/tests/opentrons/protocol_engine/execution/test_equipment_handler.py
index f740a5c3e2c..3d564f3f5e7 100644
--- a/api/tests/opentrons/protocol_engine/execution/test_equipment_handler.py
+++ b/api/tests/opentrons/protocol_engine/execution/test_equipment_handler.py
@@ -10,7 +10,7 @@
 from opentrons_shared_data.labware.dev_types import LabwareUri
 
 from opentrons.calibration_storage.helpers import uri_from_details
-from opentrons.types import Mount as HwMount, MountType, DeckSlotName, Point
+from opentrons.types import Mount as HwMount, MountType, DeckSlotName
 from opentrons.hardware_control import HardwareControlAPI
 from opentrons.hardware_control.modules import (
     TempDeck,
@@ -56,6 +56,7 @@
     LoadedPipetteData,
     LoadedModuleData,
 )
+from ..pipette_fixtures import get_default_nozzle_map
 
 
 def _make_config(use_virtual_modules: bool) -> Config:
diff --git a/api/tests/opentrons/protocol_engine/pipette_fixtures.py b/api/tests/opentrons/protocol_engine/pipette_fixtures.py
index b2d2e6bafe3..26c2ed33448 100644
--- a/api/tests/opentrons/protocol_engine/pipette_fixtures.py
+++ b/api/tests/opentrons/protocol_engine/pipette_fixtures.py
@@ -3,6 +3,9 @@
 from collections import OrderedDict
 
 from opentrons.types import Point
+from opentrons.hardware_control.nozzle_manager import NozzleMap
+from opentrons_shared_data.pipette.dev_types import PipetteNameType
+
 
 NINETY_SIX_ROWS = OrderedDict(
     (
@@ -317,3 +320,34 @@
         ("H1", Point(0.0, -31.5, 35.52)),
     )
 )
+
+
+def get_default_nozzle_map(pipette_type: PipetteNameType) -> NozzleMap:
+    """Get default nozzle map for a given pipette type."""
+    if "multi" in pipette_type.value:
+        return NozzleMap.build(
+            physical_nozzles=EIGHT_CHANNEL_MAP,
+            physical_rows=EIGHT_CHANNEL_ROWS,
+            physical_columns=EIGHT_CHANNEL_COLS,
+            starting_nozzle="A1",
+            back_left_nozzle="A1",
+            front_right_nozzle="A1",
+        )
+    elif "96" in pipette_type.value:
+        return NozzleMap.build(
+            physical_nozzles=NINETY_SIX_MAP,
+            physical_rows=NINETY_SIX_ROWS,
+            physical_columns=NINETY_SIX_COLS,
+            starting_nozzle="A1",
+            back_left_nozzle="A1",
+            front_right_nozzle="A1",
+        )
+    else:
+        return NozzleMap.build(
+            physical_nozzles=OrderedDict({"A1": Point(0, 0, 0)}),
+            physical_rows=OrderedDict({"A": ["A1"]}),
+            physical_columns=OrderedDict({"1": ["A1"]}),
+            starting_nozzle="A1",
+            back_left_nozzle="A1",
+            front_right_nozzle="A1",
+        )
diff --git a/api/tests/opentrons/protocol_engine/resources/test_pipette_data_provider.py b/api/tests/opentrons/protocol_engine/resources/test_pipette_data_provider.py
index ad0c7f5c1c0..3463c574bbc 100644
--- a/api/tests/opentrons/protocol_engine/resources/test_pipette_data_provider.py
+++ b/api/tests/opentrons/protocol_engine/resources/test_pipette_data_provider.py
@@ -1,6 +1,4 @@
 """Test pipette data provider."""
-from collections import OrderedDict
-
 import pytest
 from opentrons_shared_data.pipette.dev_types import PipetteNameType, PipetteModel
 from opentrons_shared_data.pipette import pipette_definition, types as pip_types
@@ -9,7 +7,6 @@
 )
 
 from opentrons.hardware_control.dev_types import PipetteDict
-from opentrons.hardware_control.nozzle_manager import NozzleMap
 from opentrons.protocol_engine.types import FlowRates
 from opentrons.protocol_engine.resources.pipette_data_provider import (
     LoadedStaticPipetteData,
@@ -17,7 +14,7 @@
 )
 
 from opentrons.protocol_engine.resources import pipette_data_provider as subject
-from opentrons.types import Point
+from ..pipette_fixtures import get_default_nozzle_map
 
 
 @pytest.fixture
@@ -190,6 +187,7 @@ def test_get_pipette_static_config(
     supported_tip_fixture: pipette_definition.SupportedTipsDefinition,
 ) -> None:
     """It should return config data given a PipetteDict."""
+    dummy_nozzle_map = get_default_nozzle_map(PipetteNameType.P300_SINGLE_GEN2)
     pipette_dict: PipetteDict = {
         "name": "p300_single_gen2",
         "min_volume": 20,
diff --git a/api/tests/opentrons/protocol_engine/state/test_geometry_view.py b/api/tests/opentrons/protocol_engine/state/test_geometry_view.py
index adaa65c6c3a..1268a761ff6 100644
--- a/api/tests/opentrons/protocol_engine/state/test_geometry_view.py
+++ b/api/tests/opentrons/protocol_engine/state/test_geometry_view.py
@@ -58,6 +58,7 @@
 )
 from opentrons.protocol_engine.state.addressable_areas import AddressableAreaView
 from opentrons.protocol_engine.state.geometry import GeometryView, _GripperMoveType
+from ..pipette_fixtures import get_default_nozzle_map
 
 
 @pytest.fixture
@@ -1842,6 +1843,12 @@ def test_get_next_drop_tip_location(
     decoy.when(
         labware_view.get_well_size(labware_id="abc", well_name="A1")
     ).then_return((well_size, 0, 0))
+    if pipette_channels == 96:
+        pip_type = PipetteNameType.P1000_96
+    elif pipette_channels == 8:
+        pip_type = PipetteNameType.P300_MULTI
+    else:
+        pip_type = PipetteNameType.P300_SINGLE
     decoy.when(mock_pipette_view.get_config("pip-123")).then_return(
         StaticPipetteConfig(
             min_volume=1,
diff --git a/api/tests/opentrons/protocol_engine/state/test_pipette_store.py b/api/tests/opentrons/protocol_engine/state/test_pipette_store.py
index 827df5022f5..d00599c42bd 100644
--- a/api/tests/opentrons/protocol_engine/state/test_pipette_store.py
+++ b/api/tests/opentrons/protocol_engine/state/test_pipette_store.py
@@ -52,6 +52,7 @@
     create_move_relative_command,
     create_prepare_to_aspirate_command,
 )
+from ..pipette_fixtures import get_default_nozzle_map
 
 
 @pytest.fixture
@@ -705,8 +706,8 @@ def test_add_pipette_config(
         home_position=8.9,
         nozzle_offset_z=10.11,
         bounding_nozzle_offsets=BoundingNozzlesOffsets(
-            back_left_offset=Point(x=1, y=2, z=3),
-            front_right_offset=Point(x=4, y=5, z=6),
+            back_left_offset=Point(x=0, y=0, z=0),
+            front_right_offset=Point(x=0, y=0, z=0),
         ),
         pipette_bounding_box_offsets=PipetteBoundingBoxOffsets(
             back_left_corner=Point(x=1, y=2, z=3),
diff --git a/api/tests/opentrons/protocol_engine/state/test_pipette_view.py b/api/tests/opentrons/protocol_engine/state/test_pipette_view.py
index 346d7cd5198..5de3ccd4bfc 100644
--- a/api/tests/opentrons/protocol_engine/state/test_pipette_view.py
+++ b/api/tests/opentrons/protocol_engine/state/test_pipette_view.py
@@ -38,6 +38,7 @@
     EIGHT_CHANNEL_ROWS,
     EIGHT_CHANNEL_COLS,
     EIGHT_CHANNEL_MAP,
+    get_default_nozzle_map,
 )
 
 _SAMPLE_NOZZLE_BOUNDS_OFFSETS = BoundingNozzlesOffsets(
diff --git a/api/tests/opentrons/protocol_engine/state/test_tip_state.py b/api/tests/opentrons/protocol_engine/state/test_tip_state.py
index 0bb0572e0a5..60dccf631e5 100644
--- a/api/tests/opentrons/protocol_engine/state/test_tip_state.py
+++ b/api/tests/opentrons/protocol_engine/state/test_tip_state.py
@@ -19,11 +19,12 @@
     LoadedStaticPipetteData,
 )
 from opentrons.types import Point
-
+from opentrons_shared_data.pipette.dev_types import PipetteNameType
 from ..pipette_fixtures import (
     NINETY_SIX_MAP,
     NINETY_SIX_COLS,
     NINETY_SIX_ROWS,
+    get_default_nozzle_map,
 )
 
 _tip_rack_parameters = LabwareParameters.construct(isTiprack=True)  # type: ignore[call-arg]
@@ -625,7 +626,6 @@ def test_drop_tip(
             ),
             5,
         ),
-        (None, 9),
     ],
 )
 def test_active_channels(