Reset completed iterations counter to 0 once target iterations are re…

…ached and detector is disarmed (bluesky#590)

* Reset completed iterations counter to 0 once target iterations are reached and detector is disarmed

* Fix tests

* Add test for double kickoff to increase coverage

* Make linter happy

src/ophyd_async/core/_detector.py

@@ -314,7 +314,10 @@ async def prepare(self, value: TriggerInfo) -> None:
     async def kickoff(self):
         assert self._trigger_info, "Prepare must be called before kickoff!"
         if self._iterations_completed >= self._trigger_info.iteration:
-            raise Exception(f"Kickoff called more than {self._trigger_info.iteration}")
+            raise Exception(
+                f"Kickoff called more than the configured number of "
+                f"{self._trigger_info.iteration} iteration(s)!"
+            )
         self._iterations_completed += 1
 
     @WatchableAsyncStatus.wrap
@@ -340,6 +343,7 @@ async def complete(self):
             if index >= self._trigger_info.number:
                 break
         if self._iterations_completed == self._trigger_info.iteration:
+            self._iterations_completed = 0
             await self.controller.wait_for_idle()
 
     async def describe_collect(self) -> dict[str, DataKey]:

tests/core/test_flyer.py

@@ -186,6 +186,7 @@ def flying_plan():
         yield from bps.complete(flyer, wait=False, group="complete")
         for detector in detectors:
             yield from bps.complete(detector, wait=False, group="complete")
+
         assert flyer._trigger_logic.state == TriggerState.null
 
         # Manually incremenet the index as if a frame was taken
@@ -206,6 +207,12 @@ def flying_plan():
                 name="main_stream",
             )
         yield from bps.wait(group="complete")
+
+        for detector in detectors:
+            # Since we set number of iterations to 1 (default),
+            # make sure it gets reset on complete
+            assert detector._iterations_completed == 0
+
         yield from bps.close_run()
 
         yield from bps.unstage_all(flyer, *detectors)
@@ -227,6 +234,73 @@ def flying_plan():
     ]
 
 
+async def test_hardware_triggered_flyable_too_many_kickoffs(
+    RE: RunEngine, detectors: tuple[StandardDetector]
+):
+    trigger_logic = DummyTriggerLogic()
+    flyer = StandardFlyer(trigger_logic, [], name="flyer")
+    trigger_info = TriggerInfo(
+        number=1, trigger=DetectorTrigger.constant_gate, deadtime=2, livetime=2
+    )
+
+    def flying_plan():
+        yield from bps.stage_all(*detectors, flyer)
+        assert flyer._trigger_logic.state == TriggerState.stopping
+
+        # move the flyer to the correct place, before fly scanning.
+        # Prepare the flyer first to get the trigger info for the detectors
+        yield from bps.prepare(flyer, 1, wait=True)
+
+        # prepare detectors second.
+        for detector in detectors:
+            yield from bps.prepare(
+                detector,
+                trigger_info,
+                wait=True,
+            )
+
+        yield from bps.open_run()
+        yield from bps.declare_stream(*detectors, name="main_stream", collect=True)
+
+        for _ in range(2):
+            yield from bps.kickoff(flyer)
+            for detector in detectors:
+                yield from bps.kickoff(detector)
+
+        yield from bps.complete(flyer, wait=False, group="complete")
+        for detector in detectors:
+            yield from bps.complete(detector, wait=False, group="complete")
+
+        assert flyer._trigger_logic.state == TriggerState.null
+
+        # Manually incremenet the index as if a frame was taken
+        for detector in detectors:
+            detector.writer.index += 1
+
+        yield from bps.wait(group="complete")
+
+        yield from bps.collect(
+            *detectors,
+            return_payload=False,
+            name="main_stream",
+        )
+
+        for detector in detectors:
+            # Since we set number of iterations to 1 (default),
+            # make sure it gets reset on complete
+            assert detector._iterations_completed == 0
+
+        yield from bps.close_run()
+
+        yield from bps.unstage_all(flyer, *detectors)
+
+    # fly scan
+    with pytest.raises(
+        Exception, match="Kickoff called more than the configured number"
+    ):
+        RE(flying_plan())
+
+
 # To do: Populate configuration signals
 async def test_describe_configuration():
     flyer = StandardFlyer(DummyTriggerLogic(), [], name="flyer")

tests/fastcs/panda/test_hdf_panda.py

@@ -224,10 +224,12 @@ def flying_plan():
 
         yield from bps.declare_stream(mock_hdf_panda, name="main_stream", collect=True)
 
-        for _ in range(iteration):
+        for i in range(iteration):
             set_mock_value(flyer.trigger_logic.seq.active, 1)
+
             yield from bps.kickoff(flyer, wait=True)
             yield from bps.kickoff(mock_hdf_panda)
+            assert mock_hdf_panda._iterations_completed == i + 1
 
             yield from bps.complete(flyer, wait=False, group="complete")
             yield from bps.complete(mock_hdf_panda, wait=False, group="complete")
@@ -250,6 +252,14 @@ def flying_plan():
                     name="main_stream",
                 )
             yield from bps.wait(group="complete")
+
+            # Make sure first complete doesn't reset iterations completed
+            if i == 0:
+                assert mock_hdf_panda._iterations_completed == 1
+
+        # Make sure the number of iterations completed is set to 0 after final complete.
+        assert mock_hdf_panda._iterations_completed == 0
+
         yield from bps.close_run()
 
         yield from bps.unstage_all(flyer, mock_hdf_panda)