Refactored JDS6600 tests, added amplitude implementation

1 files changed, 11 insertions, 0 deletions

diff --git a/lab_control/test/mock_jds6600_device.py b/lab_control/test/mock_jds6600_device.py
index a8f755b..c8025dd 100644
--- a/lab_control/test/mock_jds6600_device.py
+++ b/lab_control/test/mock_jds6600_device.py
@@ -9,6 +9,7 @@ class MockJDS6600Device():
         def __init__(self):
             self.on = False
             self.frequency = 0.0
+            self.amplitude = 0.0
 
     def __init__(self):
         self._master, self._slave = pty.openpty()
@@ -49,11 +50,18 @@ class MockJDS6600Device():
                 self._channels[1].on = args[1] == "1"
 
             # channel frequency
+            # TODO implement correct behaviour
             elif function == 23 or function == 24:
                 ch = function - 23
                 frequency = float(args[0])
                 self._channels[ch].frequency = frequency
 
+            # channel amplitude
+            elif function == 25 or function == 26:
+                ch = function - 25
+                amplitude = float(args[0]) / 1000.0
+                self._channels[ch].amplitude = amplitude
+
         elif opcode == "r":
             if function == 20:
                 return f":r20={int(self._channels[0].on)},{int(self._channels[1].on)}.\r\n"
@@ -85,3 +93,6 @@ class MockJDS6600Device():
     def getFrequency(self, ch: int) -> float:
         return self._channels[ch - 1].frequency
 
+    def getAmplitude(self, ch: int) -> float:
+        return self._channels[ch - 1].amplitude
+