From 5b35f32a422cc3e23deae15e332f8d0df49e6485 Mon Sep 17 00:00:00 2001
From: linuswck <linuswck@m-labs.hk>
Date: Fri, 19 Apr 2024 15:39:29 +0800
Subject: [PATCH] driver: Handle CancelledError for send_cmd handler

---
 pykirdy/driver/kirdy_async.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/pykirdy/driver/kirdy_async.py b/pykirdy/driver/kirdy_async.py
index d63e6ff..ec02e41 100644
--- a/pykirdy/driver/kirdy_async.py
+++ b/pykirdy/driver/kirdy_async.py
@@ -1,6 +1,7 @@
 import socket
 import asyncio
 import json
+import logging
 
 # Data Type Enums
 IP_SETTINGS = "ip_settings"
@@ -18,6 +19,9 @@ class CmdDoesNotExist(Exception):
 class InvalidDataType(Exception):
     pass
 
+class InvalidCmd(Exception):
+    pass
+
 class NoAckRecv(Exception):
     pass
 
@@ -661,14 +665,19 @@ class Kirdy:
     async def _send_raw_cmd(self, cmd, msg_type):
         retry = 0
         while retry < 10:
-            self._writer.write(bytes(json.dumps(cmd), "UTF-8"))
-            await self._writer.drain()
-            response = await self._read_response()
             try:
+                self._writer.write(bytes(json.dumps(cmd), "UTF-8"))
+                await self._writer.drain()
+                response = await self._read_response()
                 if response["msg_type"] == msg_type:
                     return response
-            except:
+                if response["msg_type"] == "InvalidCmd":
+                    return InvalidCmd
+            except asyncio.exceptions.CancelledError:
+                return None
+            except Exception as e:
                 retry += 1
+                logging.error(f"_send_raw_cmd Exception: {e}")
                 await asyncio.sleep(0.1)
         raise NoAckRecv
 
@@ -704,14 +713,17 @@ class Kirdy:
 
         retry = 0
         while retry < 10:
-            self._writer.write(bytes(json.dumps(cmd_dict), "UTF-8"))
-            await self._writer.drain()
-            response = await self._read_response()
             try:
+                self._writer.write(bytes(json.dumps(cmd_dict), "UTF-8"))
+                await self._writer.drain()
+                response = await self._read_response()
                 if response["msg_type"] == msg_type:
                     return response
-            except:
+            except asyncio.exceptions.CancelledError:
+                return None
+            except Exception as e:
                 retry += 1
+                logging.error(f"_send_cmd Exception: {e}")
                 await asyncio.sleep(0.1)
         raise NoAckRecv