From 32737a35e7d7d6f5aa4664c3d0711e5011ce7508 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= <gallegoj@uw.edu>
Date: Tue, 5 Dec 2023 22:38:20 -0800
Subject: [PATCH] Also record the overheads when running multiple stages
 concurrently

---
 src/hal/macros/macro.py         | 17 +++++++++++++++--
 tests/macros/test_goto_field.py | 21 ++++++++++++++++++---
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/hal/macros/macro.py b/src/hal/macros/macro.py
index e1dd65a..c3b50f5 100644
--- a/src/hal/macros/macro.py
+++ b/src/hal/macros/macro.py
@@ -415,11 +415,20 @@ async def _do_run(self):
         current_task: asyncio.Future | None = None
 
         for istage, stage in enumerate(self.stages):
+            coros = self._get_coros(stage)
+            print(coros)
             wrapped_coros = [
-                asyncio.create_task(record_overhead(self)(coro))
-                for coro in self._get_coros(stage)
+                asyncio.create_task(record_overhead(self)(coro)) for coro in coros
             ]
 
+            # If we are running multiple stages concurrently, we also record the
+            # overhead of the entire set.
+            overhead_helper: OverheadHelper | None = None
+            if len(wrapped_coros) > 1:
+                costage_name = ":".join([coro.__name__ for coro in coros])
+                overhead_helper = OverheadHelper(self, costage_name)
+                await overhead_helper.start()
+
             current_task = asyncio.gather(*wrapped_coros)
 
             self.set_stage_status(stage, StageStatus.ACTIVE)
@@ -475,6 +484,10 @@ async def _do_run(self):
                 await self.fail_macro(err, stage=stage)
                 return
 
+            finally:
+                if overhead_helper is not None:
+                    await overhead_helper.stop()
+
     def get_active_stages(self):
         """Returns a list of running stages."""
 
diff --git a/tests/macros/test_goto_field.py b/tests/macros/test_goto_field.py
index ff70efb..115d989 100644
--- a/tests/macros/test_goto_field.py
+++ b/tests/macros/test_goto_field.py
@@ -8,7 +8,12 @@
 
 from __future__ import annotations
 
+import asyncio
+
 import pytest
+from pytest_mock import MockerFixture
+
+from hal.exceptions import MacroError
 
 
 pytestmark = [pytest.mark.asyncio]
@@ -29,6 +34,7 @@ async def goto_field_macro(actor, command, mocker):
         None,
         None,
         None,
+        False,
     ]
 
     actor.models["tcc"]["axePos"].value = [100, 60, 0]
@@ -42,12 +48,21 @@ async def goto_field_macro(actor, command, mocker):
     yield macro
 
 
-async def test_goto_field_fails_tcc(goto_field_macro):
+async def test_goto_field_fails_tcc(goto_field_macro, mocker: MockerFixture):
+    mocker.patch.object(goto_field_macro, "_all_lamps_off", return_value=True)
+
+    # This causes the slew stage to fail immediately since the first thing it does
+    # is sleep for a bit.
+    mocker.patch.object(asyncio, "sleep", side_effect=MacroError)
     await goto_field_macro.run()
 
+    command = goto_field_macro.command
+
     # Macros don't finish commands.
-    assert not goto_field_macro.command.status.is_done
+    assert not command.status.is_done
     assert not goto_field_macro.running
 
-    reply_codes = [reply.flag for reply in goto_field_macro.command.actor.mock_replies]
+    reply_codes = [reply.flag for reply in command.actor.mock_replies]
     assert "e" in reply_codes
+
+    assert command.replies[-1].message["stage_duration"][1] == "slew:reconfigure"