From 8f4769503e950f88b2d36b7aa04def933814b96f Mon Sep 17 00:00:00 2001
From: ahiuchingau <20424172+ahiuchingau@users.noreply.github.com>
Date: Thu, 7 Mar 2024 14:15:50 -0500
Subject: [PATCH] broadcast request msg

---
 .../backends/ot3controller.py                 | 16 +++---
 .../hardware_control/motor_enable_disable.py  | 56 +++++++++++--------
 2 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/api/src/opentrons/hardware_control/backends/ot3controller.py b/api/src/opentrons/hardware_control/backends/ot3controller.py
index 35f6c2ba4cf..0b03cdceabe 100644
--- a/api/src/opentrons/hardware_control/backends/ot3controller.py
+++ b/api/src/opentrons/hardware_control/backends/ot3controller.py
@@ -87,7 +87,6 @@
     set_enable_tip_motor,
     set_disable_tip_motor,
     get_motor_enabled,
-    get_tip_motor_enabled,
 )
 from opentrons_hardware.hardware_control.motor_position_status import (
     get_motor_position,
@@ -1179,17 +1178,16 @@ def engaged_axes(self) -> OT3AxisMap[bool]:
     async def update_engaged_axes(self) -> None:
         """Update engaged axes."""
         motor_nodes = self._motor_nodes()
-        for node in motor_nodes:
-            await self.is_motor_engaged(node_to_axis(node))
+        results = await get_motor_enabled(self._messenger, motor_nodes)
+        for node, status in results.items():
+            self._engaged_axes[node_to_axis(node)] = status
 
     async def is_motor_engaged(self, axis: Axis) -> bool:
         node = axis_to_node(axis)
-        if axis == Axis.Q:
-            result = await get_tip_motor_enabled(self._messenger, node)
-        else:
-            result = await get_motor_enabled(self._messenger, node)
-        self._engaged_axes.update({axis: result})
-        return result
+        result = await get_motor_enabled(self._messenger, {node})
+        engaged = result[node]
+        self._engaged_axes.update({axis: engaged})
+        return engaged
 
     async def disengage_axes(self, axes: List[Axis]) -> None:
         """Disengage axes."""
diff --git a/hardware/opentrons_hardware/hardware_control/motor_enable_disable.py b/hardware/opentrons_hardware/hardware_control/motor_enable_disable.py
index d0c78b8a703..9928b841da9 100644
--- a/hardware/opentrons_hardware/hardware_control/motor_enable_disable.py
+++ b/hardware/opentrons_hardware/hardware_control/motor_enable_disable.py
@@ -1,5 +1,5 @@
 """Utilities for updating the enable/disable state of an OT3 axis."""
-from typing import Set
+from typing import Dict, Set
 import logging
 import asyncio
 from opentrons_hardware.drivers.can_bus.can_messenger import (
@@ -21,6 +21,7 @@
     ErrorCode,
     MessageId,
 )
+from opentrons_hardware.firmware_bindings.messages.messages import MessageDefinition
 
 log = logging.getLogger(__name__)
 
@@ -87,30 +88,41 @@ async def set_disable_tip_motor(
 
 async def get_motor_enabled(
     can_messenger: CanMessenger,
-    node: NodeId,
-    timeout: float = 0.5,
-) -> bool:
-    """Get motor status of a node."""
+    nodes: Set[NodeId],
+    timeout: float = 1.0,
+) -> Dict[NodeId, bool]:
+    """Get motor status of a set of nodes."""
+    expected = nodes or set()
+    reported: Dict[NodeId, bool] = {}
+    event = asyncio.Event()
 
-    def _filter(arbitration_id: ArbitrationId) -> bool:
-        return (NodeId(arbitration_id.parts.originating_node_id) == node) and (
-            MessageId(arbitration_id.parts.message_id) == GetStatusResponse.message_id
-        )
+    def _filter(arb_id: ArbitrationId) -> bool:
+        return MessageId(arb_id.parts.message_id) == GetStatusResponse.message_id
 
-    async def _wait_for_response(reader: WaitableCallback) -> bool:
+    def _listener(message: MessageDefinition, arb_id: ArbitrationId) -> None:
         """Listener for receving motor status messages."""
-        async for response, _ in reader:
-            if isinstance(response, GetStatusResponse):
-                return bool(response.payload.status.value)
-        raise StopAsyncIteration
-
-    with WaitableCallback(can_messenger, _filter) as reader:
-        await can_messenger.send(node_id=node, message=GetStatusRequest())
-        try:
-            return await asyncio.wait_for(_wait_for_response(reader), timeout)
-        except asyncio.TimeoutError:
-            log.warning("Read motor status timed out")
-            raise StopAsyncIteration
+        if isinstance(message, GetStatusResponse):
+            reported[NodeId(arb_id.parts.originating_node_id)] = bool(
+                message.payload.status.value
+            )
+
+        # found all expected nodes
+        if expected.issubset(reported):
+            event.set()
+
+    can_messenger.add_listener(_listener, _filter)
+    await can_messenger.send(node_id=NodeId.broadcast, message=GetStatusRequest())
+    try:
+        await asyncio.wait_for(event.wait(), timeout)
+    except asyncio.TimeoutError:
+        if expected:
+            log.warning(
+                "Read motor status timed out, missing nodes: "
+                f"{expected.difference(reported)}"
+            )
+        else:
+            log.debug("Read motor status terminated, no missing nodes.")
+    return reported
 
 
 async def get_tip_motor_enabled(