1 files changed, 29 insertions, 24 deletions
diff --git a/lab_control/test/mock_sds1000xe_server.py b/lab_control/test/mock_sds1000xe_server.py
index 7da90c9..5e222c7 100644
--- a/lab_control/test/mock_sds1000xe_server.py
+++ b/lab_control/test/mock_sds1000xe_server.py
@@ -1,43 +1,48 @@
 import socket
 import threading
+import atexit
 
-class MockSDS1000XEServer:
-    IP = "0.0.0.0"
-    PORT = 5025
+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.bind((IP, PORT))
+_serverSocket.listen(1)
+
+# Close it when the program exits
+atexit.register(_serverSocket.close)
 
+class MockSDS1000XEServer:
     MOTD = "Hello".encode()
 
     def __init__(self):
-        self.serverSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.serverSocket.bind((MockSDS1000XEServer.IP, MockSDS1000XEServer.PORT))
-        self.serverSocket.listen(1)
-        print("server socket is bound")
-
-        self.main = threading.Thread(target=self._mainLoop)
-        self.main.start()
+        self._stopFlag = False
+        self._clientSocket = None
+        self._mainThread = threading.Thread(target=self._mainLoop)
+        self._mainThread.start()
 
     def _mainLoop(self) -> None:
-        clientSocket, _ = self.serverSocket.accept()
-        clientSocket.settimeout(0.1)
-        clientSocket.send(MockSDS1000XEServer.MOTD)
+        self._clientSocket, _ = _serverSocket.accept()
+        self._clientSocket.settimeout(0.1)
+        self._clientSocket.send(MockSDS1000XEServer.MOTD)
 
         try:
-            while True:
-                request = clientSocket.recv(4096).decode()
+            while not self._stopFlag:
+                request = self._clientSocket.recv(4096).decode()
                 if request != "":
                     print(f"** Received request: {request}")
                     response = self._handleRequest(request)
-                    clientSocket.send(response.encode())
-        except Exception as e:
-            print("Exception caught in main server loop")
-            print(e)
+                    self._clientSocket.send(response.encode())
+        except TimeoutError as e:
+            pass
+        finally:
+            self._clientSocket.close()
     
     def _handleRequest(self, request: str) -> str:
         return "NOP"
 
     def stop(self):
-        print("Closing server socket")
-        self.serverSocket.close()
-        print("joining main")
-        self.main.join()
-        print("stopped")
+        self._stopFlag = True
+        self._mainThread.join()
+