fix: patch N1081B, adding timeout to connect()

- also update tests
ENRG-tr · Oct 16, 2024 · e9700c1 · e9700c1
1 parent c84735d
commit e9700c1
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 18 deletions.
diff --git a/src/daq/jobs/caen/n1081b.py b/src/daq/jobs/caen/n1081b.py
@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 
 from N1081B import N1081B
-from websocket import WebSocket
+from websocket import WebSocket, create_connection
 
 from daq.base import DAQJob
 from daq.models import DAQJobMessage
@@ -21,14 +21,25 @@ class DAQJobN1081BConfig(StorableDAQJobConfig):
     sections_to_store: list[str]
 
 
+class N1081BPatched(N1081B):
+    def __init__(self, ip):
+        super().__init__(ip)
+
+    def connect(self):
+        self.ws = create_connection(
+            self.API_ENDPOINT, timeout=N1081B_WEBSOCKET_TIMEOUT_SECONDS
+        )
+        return self.ws.connected
+
+
 class DAQJobN1081B(DAQJob):
     config_type = DAQJobN1081BConfig
     device: N1081B
     config: DAQJobN1081BConfig
 
     def __init__(self, config: DAQJobN1081BConfig):
         super().__init__(config)
-        self.device = N1081B(f"{config.host}:{config.port}?")
+        self.device = N1081BPatched(f"{config.host}:{config.port}?")
 
         for section in config.sections_to_store:
             if section not in N1081B.Section.__members__:

diff --git a/src/tests/test_n1081b.py b/src/tests/test_n1081b.py
@@ -1,10 +1,9 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
-from N1081B import N1081B
 from websocket import WebSocket
 
-from daq.jobs.caen.n1081b import DAQJobN1081B, DAQJobN1081BConfig
+from daq.jobs.caen.n1081b import DAQJobN1081B, DAQJobN1081BConfig, N1081BPatched
 from daq.store.models import DAQJobMessageStore
 
 
@@ -20,30 +19,30 @@ def setUp(self):
         )
         self.daq_job = DAQJobN1081B(self.config)
 
-    @patch.object(N1081B, "connect", return_value=True)
-    @patch.object(N1081B, "login", return_value=True)
+    @patch.object(N1081BPatched, "connect", return_value=True)
+    @patch.object(N1081BPatched, "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)
+    @patch.object(N1081BPatched, "connect", return_value=False)
     def test_connect_to_device_failure(self, mock_connect):
         with self.assertRaises(Exception) as context:
             self.daq_job._connect_to_device()
         self.assertTrue("Connection failed" in str(context.exception))
 
-    @patch.object(N1081B, "login", return_value=False)
-    @patch.object(N1081B, "connect", return_value=True)
+    @patch.object(N1081BPatched, "login", return_value=False)
+    @patch.object(N1081BPatched, "connect", return_value=True)
     def test_login_failure(self, mock_connect, mock_login):
         with self.assertRaises(Exception) as context:
             self.daq_job._connect_to_device()
         self.assertTrue("Login failed" in str(context.exception))
 
     @patch.object(
-        N1081B,
+        N1081BPatched,
         "get_function_results",
         return_value={"data": {"counters": [{"lemo": 1, "value": 100}]}},
     )
@@ -53,7 +52,7 @@ def test_poll_sections(self, mock_get_function_results):
         self.daq_job._send_store_message.assert_called()
         self.assertEqual(self.daq_job._send_store_message.call_count, 2)
 
-    @patch.object(N1081B, "get_function_results", return_value={"data": {}})
+    @patch.object(N1081BPatched, "get_function_results", return_value={"data": {}})
     def test_poll_sections_no_counters(self, mock_get_function_results):
         self.daq_job._send_store_message = MagicMock()
         with self.assertRaises(Exception) as context:
@@ -73,7 +72,7 @@ def test_start(
         mock_connect_to_device.assert_called_once()
         mock_poll_sections.assert_called_once()
 
-    @patch.object(N1081B, "get_function_results", return_value=None)
+    @patch.object(N1081BPatched, "get_function_results", return_value=None)
     def test_poll_sections_no_results(self, mock_get_function_results):
         self.daq_job._send_store_message = MagicMock()
         with self.assertRaises(Exception) as context:
@@ -82,7 +81,7 @@ def test_poll_sections_no_results(self, mock_get_function_results):
         self.daq_job._send_store_message.assert_not_called()
 
     @patch.object(
-        N1081B,
+        N1081BPatched,
         "get_function_results",
         return_value={"data": {"counters": [{"lemo": 1, "value": 100}]}},
     )
@@ -112,8 +111,8 @@ 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)
+    @patch.object(N1081BPatched, "connect", return_value=True)
+    @patch.object(N1081BPatched, "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"))
@@ -123,16 +122,18 @@ def test_connect_to_device_timeout(self, mock_login, mock_connect):
         mock_connect.assert_called_once()
         mock_login.assert_called_once_with("password")
 
-    @patch.object(N1081B, "get_function_results", side_effect=Exception("Timeout"))
+    @patch.object(
+        N1081BPatched, "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)
+    @patch.object(N1081BPatched, "connect", return_value=True)
+    @patch.object(N1081BPatched, "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: