From 546ae633f9aaf0b491ce4536d2987e0e6bc2bbc8 Mon Sep 17 00:00:00 2001
From: Steven B <51370195+sdb9696@users.noreply.github.com>
Date: Mon, 6 Jan 2025 12:39:46 +0000
Subject: [PATCH] Raise kasa TimeoutError instead of _RetryableError on
 asyncio.TimeoutError

---
 kasa/transports/xortransport.py     | 5 +++--
 tests/protocols/test_iotprotocol.py | 6 +++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/kasa/transports/xortransport.py b/kasa/transports/xortransport.py
index 6d04240e..8cce6eb5 100644
--- a/kasa/transports/xortransport.py
+++ b/kasa/transports/xortransport.py
@@ -23,6 +23,7 @@ from collections.abc import Generator
 
 from kasa.deviceconfig import DeviceConfig
 from kasa.exceptions import KasaException, _RetryableError
+from kasa.exceptions import TimeoutError as KasaTimeoutError
 from kasa.json import loads as json_loads
 
 from .basetransport import BaseTransport
@@ -128,7 +129,7 @@ class XorTransport(BaseTransport):
             await self._connect(self._timeout)
         except TimeoutError as ex:
             await self.reset()
-            raise _RetryableError(
+            raise KasaTimeoutError(
                 f"Timeout after {self._timeout} seconds connecting to the device:"
                 f" {self._host}:{self._port}: {ex}"
             ) from ex
@@ -167,7 +168,7 @@ class XorTransport(BaseTransport):
                 return await self._execute_send(request)
         except TimeoutError as ex:
             await self.reset()
-            raise _RetryableError(
+            raise KasaTimeoutError(
                 f"Timeout after {self._timeout} seconds sending request to the device"
                 f" {self._host}:{self._port}: {ex}"
             ) from ex
diff --git a/tests/protocols/test_iotprotocol.py b/tests/protocols/test_iotprotocol.py
index a185c70a..fd8facc9 100644
--- a/tests/protocols/test_iotprotocol.py
+++ b/tests/protocols/test_iotprotocol.py
@@ -16,7 +16,7 @@ import pytest
 from kasa.credentials import Credentials
 from kasa.device import Device
 from kasa.deviceconfig import DeviceConfig
-from kasa.exceptions import KasaException, _RetryableError
+from kasa.exceptions import KasaException, TimeoutError
 from kasa.iot import IotDevice
 from kasa.protocols.iotprotocol import IotProtocol, _deprecated_TPLinkSmartHomeProtocol
 from kasa.protocols.protocol import (
@@ -444,7 +444,7 @@ async def test_protocol_handles_timeout_failure_during_write(
     protocol = protocol_class(transport=transport_class(config=config))
     mocker.patch("asyncio.open_connection", side_effect=aio_mock_writer)
     with pytest.raises(
-        _RetryableError,
+        TimeoutError,
         match="Timeout after 5 seconds sending request to the device 127.0.0.1:9999: Simulated timeout",
     ):
         await protocol.query({})
@@ -490,7 +490,7 @@ async def test_protocol_handles_timeout_failure_during_connection(
 
     mocker.patch("asyncio.open_connection", side_effect=aio_mock_writer)
     with pytest.raises(
-        _RetryableError,
+        TimeoutError,
         match="Timeout after 5 seconds connecting to the device: 127.0.0.1:9999: Simulated timeout",
     ):
         await protocol.query({})