From 5d650761c7a007463452fb9b5ef1f9f8fe2d618c Mon Sep 17 00:00:00 2001
From: ahiuchingau <20424172+ahiuchingau@users.noreply.github.com>
Date: Thu, 9 Jan 2025 17:08:55 -0500
Subject: [PATCH] add flex stacker data as rebase

---
 .../hardware_control/modules/flex_stacker.py  | 19 +++++++++----------
 .../hardware_control/modules/types.py         | 16 ++++++++++++++++
 .../modules/module_data_mapper.py             |  1 +
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/api/src/opentrons/hardware_control/modules/flex_stacker.py b/api/src/opentrons/hardware_control/modules/flex_stacker.py
index 5ded3b391b3..18f1c5c691f 100644
--- a/api/src/opentrons/hardware_control/modules/flex_stacker.py
+++ b/api/src/opentrons/hardware_control/modules/flex_stacker.py
@@ -30,6 +30,7 @@
     StackerAxisState,
     UploadFunction,
     LiveData,
+    FlexStackerData,
 )
 
 log = logging.getLogger(__name__)
@@ -189,17 +190,15 @@ def is_simulated(self) -> bool:
 
     @property
     def live_data(self) -> LiveData:
-        return {
-            "status": self.status.value,
-            "data": {
-                "latchState": self.latch_state.value,
-                "platformState": self.platform_state.value,
-                "hopperDoorState": self.hopper_door_state.value,
-                "axisStateX": self.limit_switch_status[StackerAxis.X].value,
-                "axisStateZ": self.limit_switch_status[StackerAxis.Z].value,
-                "errorDetails": self._reader.error,
-            },
+        data: FlexStackerData = {
+            "latchState": self.latch_state.value,
+            "platformState": self.platform_state.value,
+            "hopperDoorState": self.hopper_door_state.value,
+            "axisStateX": self.limit_switch_status[StackerAxis.X].value,
+            "axisStateZ": self.limit_switch_status[StackerAxis.Z].value,
+            "errorDetails": self._reader.error,
         }
+        return {"status": self.status.value, "data": data}
 
     async def prep_for_update(self) -> str:
         await self._poller.stop()
diff --git a/api/src/opentrons/hardware_control/modules/types.py b/api/src/opentrons/hardware_control/modules/types.py
index df9c23053ed..00924ad50af 100644
--- a/api/src/opentrons/hardware_control/modules/types.py
+++ b/api/src/opentrons/hardware_control/modules/types.py
@@ -103,6 +103,15 @@ class AbsorbanceReaderData(TypedDict):
     referenceWavelength: int
 
 
+class FlexStackerData(TypedDict):
+    latchState: str
+    platformState: str
+    hopperDoorState: str
+    axisStateX: str
+    axisStateZ: str
+    errorDetails: str | None
+
+
 ModuleData = Union[
     Dict[Any, Any],  # This allows an empty dict as module data
     MagneticModuleData,
@@ -110,6 +119,7 @@ class AbsorbanceReaderData(TypedDict):
     HeaterShakerData,
     ThermocyclerData,
     AbsorbanceReaderData,
+    FlexStackerData,
 ]
 
 
@@ -144,6 +154,12 @@ def is_absorbance_reader_data(
     ) -> TypeGuard[AbsorbanceReaderData]:
         return data is not None and "uptime" in data.keys()
 
+    @classmethod
+    def is_flex_stacker_data(
+        cls, data: ModuleData | None
+    ) -> TypeGuard[FlexStackerData]:
+        return data is not None and "platformState" in data.keys()
+
 
 class LiveData(TypedDict):
     status: str
diff --git a/robot-server/robot_server/modules/module_data_mapper.py b/robot-server/robot_server/modules/module_data_mapper.py
index f2fe85c53dd..28cce33c06b 100644
--- a/robot-server/robot_server/modules/module_data_mapper.py
+++ b/robot-server/robot_server/modules/module_data_mapper.py
@@ -168,6 +168,7 @@ def map_data(
             )
         elif module_type == ModuleType.FLEX_STACKER:
             module_cls = FlexStackerModule
+            assert ModuleDataValidator.is_flex_stacker_data(live_data["data"])
             module_data = FlexStackerModuleData(
                 status=FlexStackerStatus(live_data["status"]),
                 latchState=cast(LatchState, live_data["data"].get("latchState")),