Added function setter to JDS6600

2 files changed, 26 insertions, 9 deletions

diff --git a/lab_control/test/jds6600_test.py b/lab_control/test/jds6600_test.py
index 07c77b4..c88114e 100644
--- a/lab_control/test/jds6600_test.py
+++ b/lab_control/test/jds6600_test.py
@@ -17,13 +17,13 @@ def uut(mockDevice):
     yield uut
     uut.closePort()
 
-def checkFloatParameter(testValues, writeValue, valueInMock, expectValid=True):
+def checkNumericalParameter(testValues, writeValue, valueInMock, expectValid=True):
     for ch in AVAILABLE_CHANNELS:
-        assert valueInMock(ch) == 0.0
+        assert valueInMock(ch) == None
 
         for value in testValues:
             writeValue(ch, value)
-            expectedValue = value if expectValid else 0.0
+            expectedValue = value if expectValid else None
             assert valueInMock(ch) == expectedValue
 
 def test_serialConfiguration(mockDevice):
@@ -42,13 +42,20 @@ def test_channelOnAndOff(uut, mockDevice):
         assert not mockDevice.isOn(ch)
 
 def disabled_test_setFrequency(uut, mockDevice):
-    checkFloatParameter([100.0, 100000.0, 0.0], uut.setFrequency, mockDevice.getFrequency)
+    checkNumericalParameter([100.0, 100000.0, 0.0], uut.setFrequency, mockDevice.getFrequency)
     
 def test_setInvalidFrequency(uut, mockDevice):
-    checkFloatParameter([None, -10.0, 60000000.1], uut.setFrequency, mockDevice.getFrequency, expectValid=False)
+    checkNumericalParameter([None, -10.0, 60000000.1], uut.setFrequency, mockDevice.getFrequency, expectValid=False)
 
 def test_setAmplitude(uut, mockDevice):
-    checkFloatParameter([0.1, 1.0, 10.0, 20.0, 0.0], uut.setAmplitude, mockDevice.getAmplitude)
+    checkNumericalParameter([0.1, 1.0, 10.0, 20.0, 0.0], uut.setAmplitude, mockDevice.getAmplitude)
 
 def test_setInvalidAmplitude(uut, mockDevice):
-    checkFloatParameter([None, -0.1, -10.0, 20.1], uut.setAmplitude, mockDevice.getAmplitude, expectValid=False)
+    checkNumericalParameter([None, -0.1, -10.0, 20.1], uut.setAmplitude, mockDevice.getAmplitude, expectValid=False)
+
+def test_setFunction(uut, mockDevice):
+    checkNumericalParameter(range(0, 17), uut.setFunction, mockDevice.getFunction)
+
+def test_setInvalidFunction(uut, mockDevice):
+    checkNumericalParameter([-1, -10, 17, 20, None], uut.setFunction, mockDevice.getFunction, expectValid=False)
+
diff --git a/lab_control/test/mock_jds6600_device.py b/lab_control/test/mock_jds6600_device.py
index c8025dd..a37d8a8 100644
--- a/lab_control/test/mock_jds6600_device.py
+++ b/lab_control/test/mock_jds6600_device.py
@@ -8,8 +8,9 @@ class MockJDS6600Device():
     class ChannelState:
         def __init__(self):
             self.on = False
-            self.frequency = 0.0
-            self.amplitude = 0.0
+            self.frequency = None
+            self.amplitude = None
+            self.function = None
 
     def __init__(self):
         self._master, self._slave = pty.openpty()
@@ -62,6 +63,12 @@ class MockJDS6600Device():
                 amplitude = float(args[0]) / 1000.0
                 self._channels[ch].amplitude = amplitude
 
+            # channel function shape
+            elif function == 21 or function == 22:
+                ch = function - 21
+                shape = int(args[0])
+                self._channels[ch].function = shape
+
         elif opcode == "r":
             if function == 20:
                 return f":r20={int(self._channels[0].on)},{int(self._channels[1].on)}.\r\n"
@@ -96,3 +103,6 @@ class MockJDS6600Device():
     def getAmplitude(self, ch: int) -> float:
         return self._channels[ch - 1].amplitude
 
+    def getFunction(self, ch: int) -> int:
+        return self._channels[ch - 1].function
+