From d8d4201c8c32be5917711ce502fda760c2a59ed8 Mon Sep 17 00:00:00 2001
From: Joakim Plate <elupus@ecce.se>
Date: Tue, 28 May 2024 21:36:47 +0200
Subject: [PATCH] Add retry on protocol errors

Some tv's disconnect lingering connections at write time,
leading to an exception. Retry the request once in those
cases.
---
 haphilipsjs/__init__.py | 71 ++++++++++++++++++++++++-----------------
 tests/test_v6.py        |  9 ++++++
 2 files changed, 50 insertions(+), 30 deletions(-)

diff --git a/haphilipsjs/__init__.py b/haphilipsjs/__init__.py
index dc7d87e..f2d60e7 100644
--- a/haphilipsjs/__init__.py
+++ b/haphilipsjs/__init__.py
@@ -6,6 +6,8 @@
 from urllib.parse import quote
 from secrets import token_bytes, token_hex
 from base64 import b64decode, b64encode
+from functools import wraps
+
 
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 from cryptography.hazmat.primitives.padding import PKCS7
@@ -200,6 +202,27 @@ def __init__(self, data):
 T = TypeVar("T")
 
 
+def handle_httpx_exceptions(f):
+    """Wrap up httpx exceptions in our wanted variants."""
+    @wraps(f)
+    async def wrapper(*args, **kwds):
+        try:
+            try:
+                return await f(*args, **kwds)
+            except httpx.RemoteProtocolError as err:
+                LOG.warning("%r. We retry once, could be a reused session that was closed", err)
+                return await f(*args, **kwds)
+
+        except (httpx.ConnectTimeout, httpx.ConnectError) as err:
+            raise ConnectionFailure(err) from err
+        except (httpx.ProtocolError, httpx.ReadError) as err:
+            raise ProtocolFailure(err) from err
+        except httpx.HTTPError as err:
+            raise GeneralFailure(err) from err
+
+    return wrapper
+
+
 class PhilipsTV(object):
 
     channels: ChannelsType
@@ -569,38 +592,32 @@ def _url(self, path, protocol = None):
 
         return f"{protocol}://{self._host}:{port}/{self.api_version}/{path}"
 
+    @handle_httpx_exceptions
     async def getReq(self, path, protocol = None) -> Optional[Dict]:
-        try:
-            resp = await self.session.get(self._url(path, protocol = protocol))
-            if resp.status_code == 401:
-                raise AutenticationFailure("Authenticaion failed to device")
+        resp = await self.session.get(self._url(path, protocol = protocol))
 
-            if resp.status_code != 200:
-                LOG.debug("Get failed: %s -> %d %s", path, resp.status_code, resp.text)
-                return None
+        if resp.status_code == 401:
+            raise AutenticationFailure("Authenticaion failed to device")
 
-            LOG.debug("Get succeded: %s -> %s", path, resp.text)
-            return decode_xtv_response(resp)
-        except (httpx.ConnectTimeout, httpx.ConnectError) as err:
-            raise ConnectionFailure(err) from err
-        except httpx.HTTPError as err:
-            raise GeneralFailure(err) from err
+        if resp.status_code != 200:
+            LOG.debug("Get failed: %s -> %d %s", path, resp.status_code, resp.text)
+            return None
 
+        LOG.debug("Get succeded: %s -> %s", path, resp.text)
+        return decode_xtv_response(resp)
+  
+    @handle_httpx_exceptions
     async def _getBinary(self, path: str) -> Tuple[Optional[bytes], Optional[str]]:
 
-        try:
-            resp = await self.session.get(self._url(path))
-            if resp.status_code == 401:
-                raise AutenticationFailure("Authenticaion failed to device")
+        resp = await self.session.get(self._url(path))
+        if resp.status_code == 401:
+            raise AutenticationFailure("Authenticaion failed to device")
 
-            if resp.status_code != 200:
-                return None, None
-            return resp.content, resp.headers.get("content-type")
-        except (httpx.ConnectTimeout, httpx.ConnectError) as err:
-            raise ConnectionFailure(err) from err
-        except httpx.HTTPError as err:
-            raise GeneralFailure(err) from err
+        if resp.status_code != 200:
+            return None, None
+        return resp.content, resp.headers.get("content-type")
 
+    @handle_httpx_exceptions
     async def postReq(self, path: str, data: Any, timeout=None, protocol=None) -> Optional[Dict]:
         try:
             resp = await self.session.post(self._url(path, protocol), json=data, timeout=timeout)
@@ -616,12 +633,6 @@ async def postReq(self, path: str, data: Any, timeout=None, protocol=None) -> Op
         except httpx.ReadTimeout:
             LOG.debug("Read time out on postReq", exc_info=True)
             return None
-        except (httpx.ConnectTimeout, httpx.ConnectError) as err:
-            raise ConnectionFailure(err) from err
-        except (httpx.ProtocolError, httpx.ReadError) as err:
-            raise ProtocolFailure(err) from err
-        except httpx.HTTPError as err:
-            raise GeneralFailure(err) from err
 
     async def pairRequest(
         self,
diff --git a/tests/test_v6.py b/tests/test_v6.py
index 5a6d2ad..e01f7ca 100644
--- a/tests/test_v6.py
+++ b/tests/test_v6.py
@@ -359,6 +359,15 @@ async def test_send_key_off(client_mock, param: Param):
         await client_mock.sendKey("Standby")
 
 
+async def test_send_key_retry(client_mock, param: Param):
+    route = respx.post(f"{param.base}/input/key").mock(side_effect=httpx.RemoteProtocolError)
+
+    with pytest.raises(haphilipsjs.ProtocolFailure):
+        await client_mock.sendKey("Standby")
+
+    assert route.call_count == 2
+
+
 async def test_ambilight_mode(client_mock, param):
     await client_mock.getSystem()