feat: set timeout for DAQJobN1081B websocket connection, and fix tests

src/daq/caen/n1081b.py

@@ -10,6 +10,7 @@
 from daq.store.models import DAQJobMessageStore, StorableDAQJobConfig
 
 N1081B_QUERY_INTERVAL_SECONDS = 1
+N1081B_WEBSOCKET_TIMEOUT_SECONDS = 5
 
 
 @dataclass
@@ -65,6 +66,11 @@ def _connect_to_device(self):
         if not self.device.login(self.config.password):
             raise Exception("Login failed")
 
+        if isinstance(self.device.ws, WebSocket):
+            self.device.ws.settimeout(N1081B_WEBSOCKET_TIMEOUT_SECONDS)
+        else:
+            raise Exception("Websocket not found")
+
     def _poll_sections(self):
         for section in self.config.sections_to_store:
             section = N1081B.Section[section]

src/tests/test_csv.py

@@ -35,7 +35,7 @@ def test_handle_message_new_file(
         self.store.handle_message(message)
 
         mock_add_date.assert_called_once_with("test.csv", True, None)
-        mock_open.assert_called_once_with("test.csv", "a")
+        mock_open.assert_called_once_with("test.csv", "a", newline="")
         self.assertIn("test.csv", self.store._open_csv_files)
         file = self.store._open_csv_files["test.csv"]
         self.assertEqual(len(file.write_queue), 3)  # 1 header + 2 rows
@@ -55,7 +55,7 @@ def test_handle_message_existing_file(self, mock_exists, mock_open, mock_add_dat
         self.store.handle_message(message)
 
         mock_add_date.assert_called_once_with("test.csv", True, None)
-        mock_open.assert_called_once_with("test.csv", "a")
+        mock_open.assert_called_once_with("test.csv", "a", newline="")
         self.assertIn("test.csv", self.store._open_csv_files)
         file = self.store._open_csv_files["test.csv"]
         self.assertEqual(len(file.write_queue), 2)  # 2 rows only, no header

src/tests/test_n1081b.py

@@ -2,6 +2,7 @@
 from unittest.mock import MagicMock, patch
 
 from N1081B import N1081B
+from websocket import WebSocket
 
 from daq.caen.n1081b import DAQJobN1081B, DAQJobN1081BConfig
 from daq.store.models import DAQJobMessageStore
@@ -22,9 +23,11 @@ def setUp(self):
     @patch.object(N1081B, "connect", return_value=True)
     @patch.object(N1081B, "login", return_value=True)
     def test_connect_to_device_success(self, mock_login, mock_connect):
+        self.daq_job.device.ws = MagicMock(spec=WebSocket)
         self.daq_job._connect_to_device()
         mock_connect.assert_called_once()
         mock_login.assert_called_once_with("password")
+        self.assertTrue(isinstance(self.daq_job.device.ws, WebSocket))
 
     @patch.object(N1081B, "connect", return_value=False)
     def test_connect_to_device_failure(self, mock_connect):
@@ -109,6 +112,35 @@ def test_invalid_section_in_config(self):
             DAQJobN1081B(invalid_config)
         self.assertTrue("Invalid section: INVALID_SECTION" in str(context.exception))
 
+    @patch.object(N1081B, "connect", return_value=True)
+    @patch.object(N1081B, "login", return_value=True)
+    def test_connect_to_device_timeout(self, mock_login, mock_connect):
+        self.daq_job.device.ws = MagicMock(spec=WebSocket)
+        self.daq_job.device.ws.settimeout = MagicMock(side_effect=Exception("Timeout"))
+        with self.assertRaises(Exception) as context:
+            self.daq_job._connect_to_device()
+        self.assertTrue("Timeout" in str(context.exception))
+        mock_connect.assert_called_once()
+        mock_login.assert_called_once_with("password")
+
+    @patch.object(N1081B, "get_function_results", side_effect=Exception("Timeout"))
+    def test_poll_sections_timeout(self, mock_get_function_results):
+        self.daq_job._send_store_message = MagicMock()
+        with self.assertRaises(Exception) as context:
+            self.daq_job._poll_sections()
+        self.assertTrue("Timeout" in str(context.exception))
+        self.daq_job._send_store_message.assert_not_called()
+
+    @patch.object(N1081B, "connect", return_value=True)
+    @patch.object(N1081B, "login", return_value=True)
+    def test_connect_to_device_no_websocket(self, mock_login, mock_connect):
+        self.daq_job.device.ws = None  # type: ignore
+        with self.assertRaises(Exception) as context:
+            self.daq_job._connect_to_device()
+        self.assertTrue("Websocket not found" in str(context.exception))
+        mock_connect.assert_called_once()
+        mock_login.assert_called_once_with("password")
+
 
 if __name__ == "__main__":
     unittest.main()