updated tests

Opentrons · Aug 22, 2023 · a46bda3 · a46bda3
1 parent 2ad804e
commit a46bda3
Showing 1 changed file with 59 additions and 33 deletions.
diff --git a/api/tests/opentrons/protocol_engine/state/test_geometry_view.py b/api/tests/opentrons/protocol_engine/state/test_geometry_view.py
@@ -3,7 +3,7 @@
 
 import pytest
 from decoy import Decoy
-from typing import cast, List, Tuple, Union, Optional, NamedTuple
+from typing import cast, List, Tuple, Optional, NamedTuple
 
 from opentrons_shared_data.deck.dev_types import DeckDefinitionV3
 from opentrons_shared_data.labware.dev_types import LabwareUri
@@ -1081,8 +1081,6 @@ def test_get_labware_grip_point(
     module_view: ModuleView,
     ot2_standard_deck_def: DeckDefinitionV3,
     subject: GeometryView,
-    location: Union[DeckSlotLocation, ModuleLocation],
-    expected_center_point: Point,
 ) -> None:
     """It should get the grip point of the labware at the specified location."""
     decoy.when(
@@ -1093,49 +1091,36 @@ def test_get_labware_grip_point(
         Point(x=101, y=102, z=103)
     )
     labware_center = subject.get_labware_grip_point(
-        labware_id="labware-id", location=location
+        labware_id="labware-id", location=DeckSlotLocation(slotName=DeckSlotName.SLOT_1)
     )
 
-    assert labware_center == expected_center_point
+    assert labware_center == Point(101.0, 102.0, 203)
 
 
-@pytest.mark.parametrize(
-    argnames=["location", "expected_center_point"],
-    argvalues=[
-        (OnLabwareLocation(labwareId="labware-id"), Point(5, 10, 115.0)),
-        (ModuleLocation(moduleId="module-id"), Point(111.0, 122.0, 233)),
-    ],
-)
 def test_get_labware_grip_point_on_labware(
     decoy: Decoy,
     labware_view: LabwareView,
     module_view: ModuleView,
     ot2_standard_deck_def: DeckDefinitionV3,
     subject: GeometryView,
-    location: Union[ModuleLocation, OnLabwareLocation],
-    expected_center_point: Point,
 ) -> None:
     """It should get the grip point of a labware on another labware."""
-    if isinstance(location, ModuleLocation):
-        decoy.when(module_view.get_location("module-id")).then_return(DeckSlotLocation(slotName=DeckSlotName.SLOT_4))
-    else:
-
-        decoy.when(labware_view.get(labware_id="labware-id")).then_return(
-            LoadedLabware(
-                id="labware-id",
-                loadName="above-name",
-                definitionUri="1234",
-                location=OnLabwareLocation(labwareId="below-id"),
-            )
+    decoy.when(labware_view.get(labware_id="labware-id")).then_return(
+        LoadedLabware(
+            id="labware-id",
+            loadName="above-name",
+            definitionUri="1234",
+            location=OnLabwareLocation(labwareId="below-id"),
         )
-        decoy.when(labware_view.get(labware_id="below-id")).then_return(
-            LoadedLabware(
-                id="below-id",
-                loadName="below-name",
-                definitionUri="1234",
-                location=DeckSlotLocation(slotName=DeckSlotName.SLOT_4),
-            )
+    )
+    decoy.when(labware_view.get(labware_id="below-id")).then_return(
+        LoadedLabware(
+            id="below-id",
+            loadName="below-name",
+            definitionUri="1234",
+            location=DeckSlotLocation(slotName=DeckSlotName.SLOT_4),
         )
+    )
 
     decoy.when(labware_view.get_dimensions("below-id")).then_return(
         Dimensions(x=1000, y=1001, z=11)
@@ -1155,7 +1140,48 @@ def test_get_labware_grip_point_on_labware(
         labware_id="labware-id", location=OnLabwareLocation(labwareId="below-id")
     )
 
-    assert grip_point == expected_center_point
+    assert grip_point == Point(5, 10, 115.0)
+
+
+def test_get_labware_grip_point_for_labware_on_module(
+    decoy: Decoy,
+    labware_view: LabwareView,
+    module_view: ModuleView,
+    ot2_standard_deck_def: DeckDefinitionV3,
+    subject: GeometryView,
+) -> None:
+    """It should return the grip point for labware directly on a module."""
+    decoy.when(
+        labware_view.get_grip_height_from_labware_bottom("labware-id")
+    ).then_return(500)
+    decoy.when(module_view.get_location("module-id")).then_return(
+        DeckSlotLocation(slotName=DeckSlotName.SLOT_4)
+    )
+    decoy.when(labware_view.get_deck_definition()).then_return(ot2_standard_deck_def)
+    decoy.when(
+        module_view.get_nominal_module_offset(
+            module_id="module-id", deck_type=DeckType.OT2_STANDARD
+        )
+    ).then_return(LabwareOffsetVector(x=1, y=2, z=3))
+    decoy.when(module_view.get_connected_model("module-id")).then_return(
+        ModuleModel.MAGNETIC_MODULE_V2
+    )
+    decoy.when(
+        labware_view.get_module_overlap_offsets(
+            "labware-id", ModuleModel.MAGNETIC_MODULE_V2
+        )
+    ).then_return(OverlapOffset(x=10, y=20, z=30))
+    decoy.when(module_view.get_module_calibration_offset("module-id")).then_return(
+        ModuleOffsetVector(x=100, y=200, z=300)
+    )
+    decoy.when(labware_view.get_slot_center_position(DeckSlotName.SLOT_4)).then_return(
+        Point(100, 200, 300)
+    )
+    result_grip_point = subject.get_labware_grip_point(
+        labware_id="labware-id", location=ModuleLocation(moduleId="module-id")
+    )
+
+    assert result_grip_point == Point(x=191, y=382, z=1073)
 
 
 @pytest.mark.parametrize(