Ensure exception types in our tests match websocket-client

zeek · Dec 6, 2024 · 4acc42e · 4acc42e
1 parent 627f19f
commit 4acc42e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/tests/test_controller.py b/tests/test_controller.py
@@ -139,7 +139,7 @@ def test_connect_fails_with_refused(self):
 
     def test_connect_fails_with_timeout(self):
         controller = zeekclient.controller.Controller()
-        controller.wsock.mock_connect_exc = websocket.WebSocketTimeoutError(
+        controller.wsock.mock_connect_exc = websocket.WebSocketTimeoutException(
             "connection timed out",
         )
         # Dial down attempts and waits to make this fast:
@@ -153,7 +153,7 @@ def test_connect_fails_with_timeout(self):
 
     def test_connect_fails_with_websocket_error(self):
         controller = zeekclient.controller.Controller()
-        controller.wsock.mock_connect_exc = websocket.WebSocketError("uh-oh")
+        controller.wsock.mock_connect_exc = websocket.WebSocketException("uh-oh")
         self.assertFalse(controller.connect())
         self.assertLogLines(
             "info: connecting to controller 127.0.0.1:2149",
@@ -192,7 +192,7 @@ def test_connect_fails_with_unknown_error(self):
 
     def test_handshake_fails_with_timeout(self):
         controller = zeekclient.controller.Controller()
-        controller.wsock.mock_recv_exc = websocket.WebSocketTimeoutError(
+        controller.wsock.mock_recv_exc = websocket.WebSocketTimeoutException(
             "connection timed out",
         )
         self.assertFalse(controller.connect())
@@ -278,7 +278,7 @@ def test_receive_fails_with_protocol_data_error(self):
     def test_receive_fails_with_timeout(self):
         controller = zeekclient.controller.Controller()
         self.assertTrue(controller.connect())
-        controller.wsock.mock_recv_exc = websocket.WebSocketTimeoutError(
+        controller.wsock.mock_recv_exc = websocket.WebSocketTimeoutException(
             "connection timed out",
         )
         res, msg = controller.receive()

diff --git a/tests/websocket.py b/tests/websocket.py
@@ -6,11 +6,13 @@
 import zeekclient
 
 
-class WebSocketError(Exception):
+# This typename needs to match the one in websocket-client or tests will fail.
+class WebSocketException(Exception):  # noqa: N818
     pass
 
 
-class WebSocketTimeoutError(WebSocketError):
+# This typename needs to match the one in websocket-client or tests will fail.
+class WebSocketTimeoutException(WebSocketException):  # noqa: N818
     pass