Merge pull request #7 from NitorCreations/fixes

Properly disconnect before attempting to connect again, handle error codes in responses

custom_components/extron/extron.py

@@ -1,11 +1,13 @@
 import asyncio
 import logging
+import re
 
 from asyncio import StreamReader, StreamWriter
 from asyncio.exceptions import TimeoutError
 from enum import Enum
 
 logger = logging.getLogger(__name__)
+error_regexp = re.compile("E[0-9]{2}")
 
 
 class DeviceType(Enum):
@@ -18,6 +20,14 @@ class AuthenticationError(Exception):
     pass
 
 
+class ResponseError(Exception):
+    pass
+
+
+def is_error_response(response: str) -> bool:
+    return error_regexp.match(response) is not None
+
+
 class ExtronDevice:
     def __init__(self, host: str, port: int, password: str) -> None:
         self._host = host
@@ -38,6 +48,8 @@ async def _read_until(self, phrase: str) -> str | None:
             if b.endswith(phrase.encode()):
                 return b.decode()
 
+        return None
+
     async def attempt_login(self):
         async with self._semaphore:
             await self._read_until("Password:")
@@ -60,6 +72,10 @@ async def disconnect(self):
         self._writer.close()
         await self._writer.wait_closed()
 
+    async def reconnect(self):
+        await self.disconnect()
+        await self.connect()
+
     def is_connected(self) -> bool:
         return self._connected
 
@@ -70,22 +86,26 @@ async def _run_command_internal(self, command: str):
 
             return await self._read_until("\r\n")
 
-    async def run_command(self, command: str):
+    async def run_command(self, command: str) -> str:
         try:
             response = await asyncio.wait_for(self._run_command_internal(command), timeout=3)
 
             if response is None:
                 raise RuntimeError("Command failed")
-            else:
-                return response.strip()
+
+            if is_error_response(response):
+                raise ResponseError(f"Command failed with error code {response}")
+
+            return response.strip()
         except TimeoutError:
             raise RuntimeError("Command timed out")
         except (ConnectionResetError, BrokenPipeError):
             self._connected = False
-            logger.warning("Connection seems to be broken, will attempt to reconnect")
+            raise RuntimeError("Connection was reset")
         finally:
             if not self._connected:
-                await self.connect()
+                logger.warning("Connection seems to be broken, will attempt to reconnect")
+                await self.reconnect()
 
     async def query_model_name(self):
         return await self.run_command("1I")

pyproject.toml

@@ -51,7 +51,7 @@ exclude = [
 ignore = []
 per-file-ignores = {}
 # https://docs.astral.sh/ruff/rules/
-select = ["E4", "E7", "E9", "F", "W", "N", "UP", "I"]
+select = ["E4", "E7", "E9", "F", "W", "N", "UP", "I", "RET"]
 
 # Allow unused variables when underscore-prefixed.
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"

tests/test_extron.py

@@ -0,0 +1,15 @@
+import unittest
+
+from custom_components.extron.extron import is_error_response
+
+
+class ExtronTestCase(unittest.TestCase):
+    def test_is_error_response(self):
+        self.assertTrue(is_error_response("E01"))
+        self.assertTrue(is_error_response("E74\n"))
+        self.assertTrue(is_error_response("E23\r\n  "))
+        self.assertFalse(is_error_response("0"))
+
+
+if __name__ == "__main__":
+    unittest.main()