From 47ae7d61a92a2ba4695e19705684f7b5e2f52cb3 Mon Sep 17 00:00:00 2001
From: wyzula-jan <jan.wyzula@psi.ch>
Date: Mon, 5 May 2025 12:54:14 +0200
Subject: [PATCH] refactor(waveform): added second waveform signal to the async
 waveform sim device

---
 ophyd_devices/sim/sim_waveform.py | 19 +++++++++++++++----
 tests/test_simulation.py          |  8 ++++----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/ophyd_devices/sim/sim_waveform.py b/ophyd_devices/sim/sim_waveform.py
index 3bec4cf..2095fb0 100644
--- a/ophyd_devices/sim/sim_waveform.py
+++ b/ophyd_devices/sim/sim_waveform.py
@@ -82,9 +82,16 @@ class SimWaveform(Device):
     burst = Cpt(SetableSignal, name="burst", value=1, kind=Kind.config)
 
     waveform_shape = Cpt(SetableSignal, name="waveform_shape", value=SHAPE, kind=Kind.config)
-    waveform = Cpt(
+    waveform_1 = Cpt(
         ReadOnlySignal,
-        name="waveform",
+        name="waveform_1",
+        value=np.empty(SHAPE, dtype=BIT_DEPTH),
+        compute_readback=True,
+        kind=Kind.hinted,
+    )
+    waveform_2 = Cpt(
+        ReadOnlySignal,
+        name="waveform_2",
         value=np.empty(SHAPE, dtype=BIT_DEPTH),
         compute_readback=True,
         kind=Kind.hinted,
@@ -124,6 +131,7 @@ class SimWaveform(Device):
         if self.sim_init:
             self.sim.set_init(self.sim_init)
         self._slice_index = 0
+        self.signals = [self.waveform_1, self.waveform_2]
 
     @property
     def delay_slice_update(self) -> bool:
@@ -154,7 +162,7 @@ class SimWaveform(Device):
             try:
                 for _ in range(self.burst.get()):
                     # values of the Waveform
-                    values = self.waveform.get()
+                    values = [signal.get() for signal in self.signals]
                     # add_slice option
                     if self.async_update.get() == "add_slice":
                         size = self.slice_size.get()
@@ -224,7 +232,10 @@ class SimWaveform(Device):
             )
 
         msg = messages.DeviceMessage(
-            signals={self.waveform.name: {"value": value, "timestamp": time.time()}},
+            signals={
+                self.waveform_1.name: {"value": value[0], "timestamp": time.time()},
+                self.waveform_2.name: {"value": value[1], "timestamp": time.time()},
+            },
             metadata=metadata,
         )
         # Send the message to BEC
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index 1d6390a..f17a35d 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -741,15 +741,15 @@ def test_waveform(waveform):
     """Test the SimWaveform class"""
     waveform.sim.select_model("GaussianModel")
     waveform.sim.params = {"amplitude": 500, "center": 500, "sigma": 10}
-    data = waveform.waveform.get()
+    data = waveform.waveform_1.get()
     assert isinstance(data, np.ndarray)
     assert data.shape == waveform.SHAPE
     assert np.isclose(np.argmax(data), 500, atol=5)
     waveform.waveform_shape.put(50)
-    data = waveform.waveform.get()
+    data = waveform.waveform_1.get()
     for model in waveform.sim.get_all_sim_models():
         waveform.sim.select_model(model)
-        waveform.waveform.get()
+        waveform.waveform_1.get()
     # Now also test the async readback
     mock_connector = waveform.connector = mock.MagicMock()
     mock_run_subs = waveform._run_subs = mock.MagicMock()
@@ -788,7 +788,7 @@ def test_waveform_update_modes(waveform, mode, mock_data, expected_calls):
     with (
         mock.patch.object(waveform, "_run_subs") as mock_run_subs,
         mock.patch.object(waveform, "_send_async_update") as mock_send_async_update,
-        mock.patch.object(waveform.waveform, "get", return_value=mock_data),
+        mock.patch.object(waveform.waveform_1, "get", return_value=mock_data),
     ):
 
         status = waveform.trigger()