From b66347116f3d535447c89b722672fba3347886f8 Mon Sep 17 00:00:00 2001
From: sdb9696 <51370195+sdb9696@users.noreply.github.com>
Date: Wed, 20 Dec 2023 19:16:23 +0000
Subject: [PATCH] Add optional error code to exceptions (#585)

---
 kasa/aestransport.py  |  8 ++++----
 kasa/exceptions.py    | 14 ++++++++++++++
 kasa/smartprotocol.py |  8 ++++----
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/kasa/aestransport.py b/kasa/aestransport.py
index e79d0651..e7dd5356 100644
--- a/kasa/aestransport.py
+++ b/kasa/aestransport.py
@@ -130,14 +130,14 @@ class AesTransport(BaseTransport):
             return
         msg = f"{msg}: {self._host}: {error_code.name}({error_code.value})"
         if error_code in SMART_TIMEOUT_ERRORS:
-            raise TimeoutException(msg)
+            raise TimeoutException(msg, error_code=error_code)
         if error_code in SMART_RETRYABLE_ERRORS:
-            raise RetryableException(msg)
+            raise RetryableException(msg, error_code=error_code)
         if error_code in SMART_AUTHENTICATION_ERRORS:
             self._handshake_done = False
             self._login_token = None
-            raise AuthenticationException(msg)
-        raise SmartDeviceException(msg)
+            raise AuthenticationException(msg, error_code=error_code)
+        raise SmartDeviceException(msg, error_code=error_code)
 
     async def send_secure_passthrough(self, request: str):
         """Send encrypted message as passthrough."""
diff --git a/kasa/exceptions.py b/kasa/exceptions.py
index e83c9237..ff91a7b0 100644
--- a/kasa/exceptions.py
+++ b/kasa/exceptions.py
@@ -1,10 +1,15 @@
 """python-kasa exceptions."""
 from enum import IntEnum
+from typing import Optional
 
 
 class SmartDeviceException(Exception):
     """Base exception for device errors."""
 
+    def __init__(self, *args, error_code: Optional["SmartErrorCode"] = None):
+        self.error_code = error_code
+        super().__init__(args)
+
 
 class UnsupportedDeviceException(SmartDeviceException):
     """Exception for trying to connect to unsupported devices."""
@@ -17,14 +22,23 @@ class UnsupportedDeviceException(SmartDeviceException):
 class AuthenticationException(SmartDeviceException):
     """Base exception for device authentication errors."""
 
+    def __init__(self, *args, error_code: Optional["SmartErrorCode"] = None):
+        super().__init__(args, error_code)
+
 
 class RetryableException(SmartDeviceException):
     """Retryable exception for device errors."""
 
+    def __init__(self, *args, error_code: Optional["SmartErrorCode"] = None):
+        super().__init__(args, error_code)
+
 
 class TimeoutException(SmartDeviceException):
     """Timeout exception for device errors."""
 
+    def __init__(self, *args, error_code: Optional["SmartErrorCode"] = None):
+        super().__init__(args, error_code)
+
 
 class SmartErrorCode(IntEnum):
     """Enum for SMART Error Codes."""
diff --git a/kasa/smartprotocol.py b/kasa/smartprotocol.py
index a9266174..443d1def 100644
--- a/kasa/smartprotocol.py
+++ b/kasa/smartprotocol.py
@@ -186,12 +186,12 @@ class SmartProtocol(TPLinkProtocol):
         if method := resp_dict.get("method"):
             msg += f" for method: {method}"
         if error_code in SMART_TIMEOUT_ERRORS:
-            raise TimeoutException(msg)
+            raise TimeoutException(msg, error_code=error_code)
         if error_code in SMART_RETRYABLE_ERRORS:
-            raise RetryableException(msg)
+            raise RetryableException(msg, error_code=error_code)
         if error_code in SMART_AUTHENTICATION_ERRORS:
-            raise AuthenticationException(msg)
-        raise SmartDeviceException(msg)
+            raise AuthenticationException(msg, error_code=error_code)
+        raise SmartDeviceException(msg, error_code=error_code)
 
     async def close(self) -> None:
         """Close the protocol."""