Catch websocket errors within the library and raise Library own ValloxApiException or ValloxWebsocketException

Author: yozik04

setup.py

@@ -35,7 +35,7 @@ def parse_requirements(filename, parent=None):
 setup(
   name='vallox_websocket_api',
   packages=['vallox_websocket_api'],
-  version='2.0.0',
+  version='2.1.0',
   python_requires=">=3.5.1, <4",
   description='Vallox WebSocket API',
   author='Jevgeni Kiski',

tests/test_client.py

@@ -2,6 +2,11 @@
 import asynctest
 
 from tests.decorators import with_client
+from asynctest import CoroutineMock
+
+from websockets.exceptions import InvalidMessage
+
+from vallox_websocket_api.exceptions import ValloxWebsocketException
 
 class TestClient(asynctest.TestCase):
   @with_client
@@ -87,4 +92,14 @@ async def testSetNewSettableAddressByAddress(self, client, ws):
 
     await client.set_values({
       'A_CYC_RH_VALUE': 22
-    })
+    })
+
+  @with_client
+  async def testSetNewSettableAddressByAddressException(self, client, ws):
+    ws.recv.side_effect = CoroutineMock(side_effect=InvalidMessage())
+
+    client.set_settable_address(4363, int)
+
+    await self.assertAsyncRaisesRegex(ValloxWebsocketException, "Websocket handshake failed", client.set_values({
+      'A_CYC_RH_VALUE': 22
+    }))

vallox_websocket_api/__init__.py

@@ -1,4 +1,5 @@
 from .client import Client
 from .vallox import Vallox, PROFILE
+from .exceptions import ValloxApiException, ValloxWebsocketException
 
-__all__ = ["Client", "Vallox", "PROFILE"]
+__all__ = ["Client", "Vallox", "PROFILE", "ValloxApiException", "ValloxWebsocketException"]

vallox_websocket_api/client.py

@@ -1,10 +1,11 @@
-import numpy
-import re
 import datetime
+import re
 
+import numpy
 import websockets
 
 from .constants import vlxDevConstants, vlxOffsetObject
+from .exceptions import ValloxWebsocketException
 
 KPageSize = 65536
 
@@ -220,15 +221,26 @@ def _decode_dict(self, dict):
     return new_dict
 
   async def _websocket_request(self, command=vlxDevConstants.WS_WEB_UI_COMMAND_WRITE_DATA, dict=None, read_packets=1):
-    async with websockets.connect("ws://%s/" % self.ip_address) as ws:
-        request = self._make_payload(command, dict)
-        await ws.send(request)
-        results = []
-        for i in range(0, read_packets):
-          r = await ws.recv()
-
-          results.append(r)
-        return results[0] if read_packets == 1 else results
+    try:
+      async with websockets.connect("ws://%s/" % self.ip_address) as ws:
+          request = self._make_payload(command, dict)
+          await ws.send(request)
+          results = []
+          for i in range(0, read_packets):
+            r = await ws.recv()
+
+            results.append(r)
+          return results[0] if read_packets == 1 else results
+    except websockets.InvalidHandshake as e:
+      raise ValloxWebsocketException('Websocket handshake failed') from e
+    except websockets.InvalidURI as e:
+      raise ValloxWebsocketException('Websocket invalid URI') from e
+    except websockets.PayloadTooBig as e:
+      raise ValloxWebsocketException('Websocket payload too big') from e
+    except websockets.InvalidState as e:
+      raise ValloxWebsocketException('Websocket invalid state') from e
+    except websockets.WebSocketProtocolError as e:
+      raise ValloxWebsocketException('Websocket protocol error') from e
 
   async def fetch_metrics(self, metric_keys=None):
     metrics = {}

vallox_websocket_api/exceptions.py

@@ -0,0 +1,7 @@
+class ValloxApiException(Exception):
+    pass
+
+class ValloxWebsocketException(ValloxApiException):
+    """
+    Is raised on any websocket error
+    """