From f68acb60e7f4ca68be2d3c40cb780d9d0b168874 Mon Sep 17 00:00:00 2001
From: sdb9696 <steven.beth@gmail.com>
Date: Tue, 23 Jan 2024 17:51:18 +0000
Subject: [PATCH] Avoid reset for all connection errors

---
 kasa/exceptions.py              | 4 ----
 kasa/httpclient.py              | 9 ++-------
 kasa/iotprotocol.py             | 7 -------
 kasa/smartprotocol.py           | 7 -------
 kasa/tests/test_httpclient.py   | 7 +++----
 kasa/tests/test_klapprotocol.py | 2 +-
 6 files changed, 6 insertions(+), 30 deletions(-)

diff --git a/kasa/exceptions.py b/kasa/exceptions.py
index 8720d97b..c0ef23b6 100644
--- a/kasa/exceptions.py
+++ b/kasa/exceptions.py
@@ -42,10 +42,6 @@ class ConnectionException(SmartDeviceException):
     """Connection exception for device errors."""
 
 
-class DisconnectedException(SmartDeviceException):
-    """Disconnected exception for device errors."""
-
-
 class SmartErrorCode(IntEnum):
     """Enum for SMART Error Codes."""
 
diff --git a/kasa/httpclient.py b/kasa/httpclient.py
index 73c91fa4..7fe0b2c3 100644
--- a/kasa/httpclient.py
+++ b/kasa/httpclient.py
@@ -7,7 +7,6 @@ import aiohttp
 from .deviceconfig import DeviceConfig
 from .exceptions import (
     ConnectionException,
-    DisconnectedException,
     SmartDeviceException,
     TimeoutException,
 )
@@ -81,13 +80,9 @@ class HttpClient:
                     if return_json:
                         response_data = json_loads(response_data.decode())
 
-        except aiohttp.ServerDisconnectedError as ex:
-            raise DisconnectedException(
-                f"Disconnected from the device: {self._config.host}: {ex}", ex
-            ) from ex
-        except aiohttp.ClientOSError as ex:
+        except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as ex:
             raise ConnectionException(
-                f"Unable to connect to the device: {self._config.host}: {ex}", ex
+                f"Device connection error: {self._config.host}: {ex}", ex
             ) from ex
         except (aiohttp.ServerTimeoutError, asyncio.TimeoutError) as ex:
             raise TimeoutException(
diff --git a/kasa/iotprotocol.py b/kasa/iotprotocol.py
index aac21f10..ed926101 100755
--- a/kasa/iotprotocol.py
+++ b/kasa/iotprotocol.py
@@ -6,7 +6,6 @@ from typing import Dict, Union
 from .exceptions import (
     AuthenticationException,
     ConnectionException,
-    DisconnectedException,
     RetryableException,
     SmartDeviceException,
     TimeoutException,
@@ -45,13 +44,7 @@ class IotProtocol(BaseProtocol):
         for retry in range(retry_count + 1):
             try:
                 return await self._execute_query(request, retry)
-            except DisconnectedException as sdex:
-                if retry >= retry_count:
-                    _LOGGER.debug("Giving up on %s after %s retries", self._host, retry)
-                    raise sdex
-                continue
             except ConnectionException as sdex:
-                await self._transport.reset()
                 if retry >= retry_count:
                     _LOGGER.debug("Giving up on %s after %s retries", self._host, retry)
                     raise sdex
diff --git a/kasa/smartprotocol.py b/kasa/smartprotocol.py
index 8b876c14..6f0648ea 100644
--- a/kasa/smartprotocol.py
+++ b/kasa/smartprotocol.py
@@ -18,7 +18,6 @@ from .exceptions import (
     SMART_TIMEOUT_ERRORS,
     AuthenticationException,
     ConnectionException,
-    DisconnectedException,
     RetryableException,
     SmartDeviceException,
     SmartErrorCode,
@@ -66,13 +65,7 @@ class SmartProtocol(BaseProtocol):
         for retry in range(retry_count + 1):
             try:
                 return await self._execute_query(request, retry)
-            except DisconnectedException as sdex:
-                if retry >= retry_count:
-                    _LOGGER.debug("Giving up on %s after %s retries", self._host, retry)
-                    raise sdex
-                continue
             except ConnectionException as sdex:
-                await self._transport.reset()
                 if retry >= retry_count:
                     _LOGGER.debug("Giving up on %s after %s retries", self._host, retry)
                     raise sdex
diff --git a/kasa/tests/test_httpclient.py b/kasa/tests/test_httpclient.py
index bcf48df2..e178b818 100644
--- a/kasa/tests/test_httpclient.py
+++ b/kasa/tests/test_httpclient.py
@@ -7,7 +7,6 @@ import pytest
 from ..deviceconfig import DeviceConfig
 from ..exceptions import (
     ConnectionException,
-    DisconnectedException,
     SmartDeviceException,
     TimeoutException,
 )
@@ -19,13 +18,13 @@ from ..httpclient import HttpClient
     [
         (
             aiohttp.ServerDisconnectedError(),
-            DisconnectedException,
-            "Disconnected from the device: ",
+            ConnectionException,
+            "Device connection error: ",
         ),
         (
             aiohttp.ClientOSError(),
             ConnectionException,
-            "Unable to connect to the device: ",
+            "Device connection error: ",
         ),
         (
             aiohttp.ServerTimeoutError(),
diff --git a/kasa/tests/test_klapprotocol.py b/kasa/tests/test_klapprotocol.py
index 6bd142ca..09ceccae 100644
--- a/kasa/tests/test_klapprotocol.py
+++ b/kasa/tests/test_klapprotocol.py
@@ -57,7 +57,7 @@ class _mock_response:
         (aiohttp.ServerDisconnectedError("dummy exception"), True),
         (aiohttp.ClientOSError("dummy exception"), True),
     ],
-    ids=("Exception", "SmartDeviceException", "DisconnectError", "ConnectError"),
+    ids=("Exception", "ServerTimeoutError", "ServerDisconnectedError", "ClientOSError"),
 )
 @pytest.mark.parametrize("transport_class", [AesTransport, KlapTransport])
 @pytest.mark.parametrize("protocol_class", [IotProtocol, SmartProtocol])