1 files changed, 37 insertions, 13 deletions
diff --git a/lab_control/test/mock_sds1000xe_server.py b/lab_control/test/mock_sds1000xe_server.py
index 5e222c7..fc6128e 100644
--- a/lab_control/test/mock_sds1000xe_server.py
+++ b/lab_control/test/mock_sds1000xe_server.py
@@ -1,48 +1,72 @@
 import socket
 import threading
 import atexit
+import re
 
 IP = "0.0.0.0"
 PORT = 5025
 
 # Bind server socket when this module is included
 _serverSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+_serverSocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 _serverSocket.bind((IP, PORT))
 _serverSocket.listen(1)
 
 # Close it when the program exits
-atexit.register(_serverSocket.close)
+def _cleanup():
+    _serverSocket.close()
+atexit.register(_cleanup)
 
 class MockSDS1000XEServer:
-    MOTD = "Hello".encode()
-
+    class ChannelValues:
+        def __init__(self):
+            self.amplitude = None
+    
     def __init__(self):
         self._stopFlag = False
         self._clientSocket = None
         self._mainThread = threading.Thread(target=self._mainLoop)
         self._mainThread.start()
 
+        # Mock measured values
+        self._channels = [MockSDS1000XEServer.ChannelValues() for i in range(0, 4)]
+
     def _mainLoop(self) -> None:
         self._clientSocket, _ = _serverSocket.accept()
         self._clientSocket.settimeout(0.1)
-        self._clientSocket.send(MockSDS1000XEServer.MOTD)
 
         try:
             while not self._stopFlag:
-                request = self._clientSocket.recv(4096).decode()
-                if request != "":
-                    print(f"** Received request: {request}")
+                try:
+                    request = self._clientSocket.recv(4096).decode()
+
                     response = self._handleRequest(request)
-                    self._clientSocket.send(response.encode())
-        except TimeoutError as e:
-            pass
+                    if response is not None:
+                        self._clientSocket.send(response.encode())
+                except TimeoutError as e:
+                    pass
         finally:
             self._clientSocket.close()
     
     def _handleRequest(self, request: str) -> str:
-        return "NOP"
+        m = re.search(r"C(?P<channel>\d):(?P<opcode>\w{4})\?\s(?P<arg>\w+)", request)
+        if not m:
+            return None
+        
+        channelIndex = int(m.group("channel")) - 1
+        if m.group("opcode") == "PAVA":
+            if m.group("arg") == "AMPL":
+                if self._channels[channelIndex].amplitude is None:
+                    return None
+                else:
+                    value = self._channels[channelIndex].amplitude
+                    response = f"C{m.group('channel')}:PAVA AMPL,{value:.6E}V"
+                    return response
 
-    def stop(self):
+    def stop(self) -> None:
         self._stopFlag = True
         self._mainThread.join()
-    
+
+    def setAmplitude(self, channel: int, ampl: float) -> None:
+        self._channels[channel - 1].amplitude = ampl
+