fix tests

flxdot · Apr 24, 2024 · 3eeac27 · 3eeac27
1 parent d38d009
commit 3eeac27
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/lib/py_edge_device/carlos/edge/device/communication.py b/lib/py_edge_device/carlos/edge/device/communication.py
@@ -4,6 +4,7 @@
     CarlosMessage,
     DeviceId,
     EdgeCommunicationHandler,
+    EdgeConnectionDisconnected,
     EdgeProtocol,
     EdgeVersionPayload,
     MessageType,
@@ -32,6 +33,16 @@ def __init__(self, protocol: EdgeProtocol, device_id: DeviceId):
             }
         )
 
+    async def listen(self):  # pragma: no cover
+        """Wrapper around the actual listen method to catch exceptions and automatically
+        reconnect."""
+
+        while not self._stopped:
+            try:
+                await super().listen()
+            except EdgeConnectionDisconnected:
+                await self.protocol.connect()
+
 
 async def handle_edge_version(
     protocol: EdgeProtocol, message: CarlosMessage

diff --git a/lib/py_edge_interface/carlos/edge/interface/protocol.py b/lib/py_edge_interface/carlos/edge/interface/protocol.py
@@ -151,11 +151,8 @@ async def listen(self):
             await self.protocol.connect()  # pragma: no cover
 
         while not self._stopped:
-            try:
-                await self.handle_message(await self.protocol.receive())
-                await sleep(0.1)
-            except EdgeConnectionDisconnected:
-                await self.protocol.connect()
+            await self.handle_message(await self.protocol.receive())
+            await sleep(0.1)
 
     async def handle_message(self, message: CarlosMessage):
         """Handles the incoming message.

diff --git a/services/api/carlos/api/routes/device_server_routes/device_server_routes.py b/services/api/carlos/api/routes/device_server_routes/device_server_routes.py
@@ -23,9 +23,9 @@
 
 from carlos.api.depends.authentication import verify_token as auth_verify_token
 from carlos.api.depends.context import request_context
+from carlos.api.depends.database import carlos_db_connection
 from carlos.api.routes.devices_routes import DEVICE_ID_PATH
 
-from ...depends.database import carlos_db_connection
 from .protocol import WebsocketProtocol
 from .state import DEVICE_CONNECTION_MANAGER