From b2ae87598e46badecb35fd721a74041138651613 Mon Sep 17 00:00:00 2001
From: appel_c <christian.appel@psi.ch>
Date: Wed, 25 Jun 2025 18:07:16 +0200
Subject: [PATCH] fix(mock_pv): add callback mechanism for signal.subscribe to
 MockPv

---
 ophyd_devices/tests/utils.py  | 13 ++++++++++---
 tests/test_delay_generator.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/ophyd_devices/tests/utils.py b/ophyd_devices/tests/utils.py
index a9940c8..f36bf07 100644
--- a/ophyd_devices/tests/utils.py
+++ b/ophyd_devices/tests/utils.py
@@ -229,13 +229,20 @@ class MockPV:
     ):
         """MOCK PV, put function"""
         self.mock_data = value
-        if callback is not None:
-            callback()
+        for cb, kw in self.callbacks.values():
+            if callable(cb):
+                cb(value, **kw)
 
     # pylint: disable=unused-argument
     def add_callback(self, callback=None, index=None, run_now=False, with_ctrlvars=True, **kw):
         """Add callback"""
-        return mock.MagicMock()
+        if index is None:
+            index = len(self.callbacks)
+        if index in self.callbacks:
+            index = len(self.callbacks)
+        if callback is not None:
+            self.callbacks[index] = (callback, kw)
+        return index
 
     # pylint: disable=unused-argument
     def get_with_metadata(
diff --git a/tests/test_delay_generator.py b/tests/test_delay_generator.py
index ca441b6..418d714 100644
--- a/tests/test_delay_generator.py
+++ b/tests/test_delay_generator.py
@@ -73,3 +73,32 @@ def test_check_if_ddg_okay(mock_ddg):
     mock_ddg.status._read_pv.mock_data = "STATUS NOT OK"
     with pytest.raises(DelayGeneratorError):
         mock_ddg.check_if_ddg_okay()
+
+
+def test_subscribe_signal(mock_ddg):
+    """Test that MockPV allows subscription to signals."""
+    arg_buffer = []
+    kwargs_buffer = []
+
+    def cb(*args, **kwargs):
+        arg_buffer.append(args)
+        kwargs_buffer.append(kwargs)
+
+    mock_ddg: DelayGenerator
+    mock_ddg.burstMode.put(0)
+    mock_ddg.burstMode._read_changed(0)
+    assert mock_ddg.burstMode.get() == 0
+    mock_ddg.burstMode.subscribe(callback=cb, run=False)
+
+    assert len(arg_buffer) == 0
+    assert len(kwargs_buffer) == 0
+    mock_ddg.burstMode.put(1)
+    assert len(arg_buffer) == 1
+    kwargs = kwargs_buffer[0]
+    assert len(kwargs_buffer) == 1
+    assert kwargs["old_value"] == 0
+    assert kwargs["value"] == 1
+    assert kwargs["timestamp"] == mock.ANY
+    assert kwargs["sub_type"] == "value"
+    assert kwargs["obj"].name == "ddg_burstMode"
+    assert isinstance(kwargs["obj"], ophyd.EpicsSignal)