Partial implementation of setFrequency, need to experiment with device to understand the scaling

3 files changed, 49 insertions, 8 deletions

diff --git a/lab_control/jds6600.py b/lab_control/jds6600.py
index 7e7640d..6137105 100644
--- a/lab_control/jds6600.py
+++ b/lab_control/jds6600.py
@@ -20,18 +20,23 @@ class JDS6600(FunctionGenerator):
     def closePort(self) -> None:
         self._port.close()
 
-    def _queryOnOff(self) -> list[int, int]:
+    def _queryOnOff(self) -> list[str, str]:
+        # TODO double check query on the device
         response = self._sendRequest("r20", "")
         return [response[5], response[7]]
 
     def setOn(self, channel: int) -> None:
         state = self._queryOnOff()
-        state[channel - 1] = '1'
+        state[channel - 1] = "1"
         response = self._sendRequest("w20", f"{state[0]},{state[1]}")
         # TODO figure out error handling
 
     def setOff(self, channel: int) -> None:
         state = self._queryOnOff()
-        state[channel - 1] = '0'
+        state[channel - 1] = "0"
         response = self._sendRequest("w20", f"{state[0]},{state[1]}")
         # TODO figure out error handling
+
+    def setFrequency(self, channel: int, frequency: float) -> None:
+        opcode = f"w{channel + 23}"
+
diff --git a/lab_control/test/jds6600_test.py b/lab_control/test/jds6600_test.py
index 16deeb4..948b5fc 100644
--- a/lab_control/test/jds6600_test.py
+++ b/lab_control/test/jds6600_test.py
@@ -31,3 +31,23 @@ def test_channelOnAndOff(uut, mockDevice):
         assert mockDevice.isOn(ch)
         uut.setOff(ch)
         assert not mockDevice.isOn(ch)
+
+def test_setFrequency(uut, mockDevice):
+    for ch in AVAILABLE_CHANNELS:
+        assert mockDevice.getFrequency(ch) == 0.0
+
+        for frequency in [100.0, 100000.0, 0.0]:
+            uut.setFrequency(ch, frequency)
+            assert mockDevice.getFrequency(ch) == frequency
+
+def test_setInvalidFrequency(uut, mockDevice):
+    for ch in AVAILABLE_CHANNELS:
+        uut.setFrequency(ch, None)
+        assert mockDevice.getFrequency(ch) == 0.0
+
+        uut.setFrequency(ch, -10.0)
+        assert mockDevice.getFrequency(ch) == 0.0
+
+        uut.setFrequency(ch, 60000000.1)
+        assert mockDevice.getFrequency(ch) == 0.0
+
diff --git a/lab_control/test/mock_jds6600_device.py b/lab_control/test/mock_jds6600_device.py
index 6e036d5..a8f755b 100644
--- a/lab_control/test/mock_jds6600_device.py
+++ b/lab_control/test/mock_jds6600_device.py
@@ -5,12 +5,17 @@ import threading
 import re
 
 class MockJDS6600Device():
+    class ChannelState:
+        def __init__(self):
+            self.on = False
+            self.frequency = 0.0
+
     def __init__(self):
         self._master, self._slave = pty.openpty()
         self._masterFile = os.fdopen(self._master, mode="r+b", closefd=False, buffering=0)
 
         self._portName = os.ttyname(self._slave)
-        self._channels = [False, False]
+        self._channels = [MockJDS6600Device.ChannelState() for i in [1, 2]]
         
         self._mainThread = threading.Thread(target=self._mainLoop)
         self._mainThread.start()
@@ -38,12 +43,20 @@ class MockJDS6600Device():
         args = m.group("args").split(",")
 
         if opcode == "w":
+            # channel on/off
             if function == 20:
-                self._channels[0] = args[0] == "1"
-                self._channels[1] = args[1] == "1"
+                self._channels[0].on = args[0] == "1"
+                self._channels[1].on = args[1] == "1"
+
+            # channel frequency
+            elif function == 23 or function == 24:
+                ch = function - 23
+                frequency = float(args[0])
+                self._channels[ch].frequency = frequency
+
         elif opcode == "r":
             if function == 20:
-                return f":r20={int(self._channels[0])},{int(self._channels[1])}.\r\n"
+                return f":r20={int(self._channels[0].on)},{int(self._channels[1].on)}.\r\n"
 
         return ":ok\r\n"
 
@@ -67,5 +80,8 @@ class MockJDS6600Device():
         return self._portName
 
     def isOn(self, ch: int) -> bool:
-        return self._channels[ch - 1]
+        return self._channels[ch - 1].on
+
+    def getFrequency(self, ch: int) -> float:
+        return self._channels[ch - 1].frequency