Add ability to disable and enable retry logic (#296)

* Add ability to disable and enable retry logic

* Typing

* Linting

Author: bachya

simplipy/api.py

@@ -5,7 +5,7 @@
 from datetime import datetime, timedelta
 from json.decoder import JSONDecodeError
 import sys
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, cast
 
 from aiohttp import ClientSession
 from aiohttp.client_exceptions import ClientResponseError
@@ -68,6 +68,7 @@ def __init__(
     ) -> None:
         """Initialize."""
         self._refresh_token_callbacks: list[Callable[..., None]] = []
+        self._request_retries = request_retries
         self.session: ClientSession = session
 
         # These will get filled in after initial authentication:
@@ -79,16 +80,7 @@ def __init__(
         self.user_id: int | None = None
         self.websocket: WebsocketClient | None = None
 
-        # Implement a version of the request coroutine, but with backoff/retry logic:
-        self.async_request = backoff.on_exception(
-            backoff.expo,
-            ClientResponseError,
-            jitter=backoff.random_jitter,
-            logger=LOGGER,
-            max_tries=request_retries,
-            on_backoff=self._async_handle_on_backoff,
-            on_giveup=self._handle_on_giveup,
-        )(self._async_request)
+        self.async_request = self._wrap_request_method(self._request_retries)
 
     @classmethod
     async def async_from_auth(
@@ -264,6 +256,29 @@ def _handle_on_giveup(_: dict[str, Any]) -> None:
         err = err_info[1].with_traceback(err_info[2])  # type: ignore
         raise RequestError(err) from err
 
+    def _wrap_request_method(self, request_retries: int) -> Callable:
+        """Wrap the request method in backoff/retry logic."""
+        return cast(
+            Callable,
+            backoff.on_exception(
+                backoff.expo,
+                ClientResponseError,
+                jitter=backoff.random_jitter,
+                logger=LOGGER,
+                max_tries=request_retries,
+                on_backoff=self._async_handle_on_backoff,
+                on_giveup=self._handle_on_giveup,
+            )(self._async_request),
+        )
+
+    def disable_request_retries(self) -> None:
+        """Disable the request retry mechanism."""
+        self.async_request = self._wrap_request_method(1)
+
+    def enable_request_retries(self) -> None:
+        """Enable the request retry mechanism."""
+        self.async_request = self._wrap_request_method(self._request_retries)
+
     def add_refresh_token_callback(
         self, callback: Callable[..., None]
     ) -> Callable[..., None]:

tests/test_api.py

@@ -278,45 +278,20 @@ async def test_refresh_token_callback(
 
 
 @pytest.mark.asyncio
-async def test_request_error_failed_retry(aresponses, server):
-    """Test that a RequestError that fails multiple times still raises."""
-    server.add(
-        "api.simplisafe.com",
-        f"/v1/users/{TEST_SUBSCRIPTION_ID}/subscriptions",
-        "get",
-        response=aresponses.Response(text="Conflict", status=409),
-    )
-    server.add(
-        "api.simplisafe.com",
-        f"/v1/users/{TEST_SUBSCRIPTION_ID}/subscriptions",
-        "get",
-        response=aresponses.Response(text="Conflict", status=409),
-    )
-
-    async with aiohttp.ClientSession() as session:
-        simplisafe = await API.async_from_auth(
-            TEST_AUTHORIZATION_CODE,
-            TEST_CODE_VERIFIER,
-            session=session,
-            # Set so that our tests don't take too long:
-            request_retries=1,
-        )
-
-        with pytest.raises(RequestError):
-            await simplisafe.async_get_systems()
-
-    aresponses.assert_plan_strictly_followed()
-
-
-@pytest.mark.asyncio
-async def test_request_error_successful_retry(
+async def test_request_retry(
     api_token_response,
     aresponses,
     server,
     v2_settings_response,
     v2_subscriptions_response,
 ):
-    """Test that a RequestError can be successfully retried."""
+    """Test that request retries work."""
+    server.add(
+        "api.simplisafe.com",
+        f"/v1/users/{TEST_SUBSCRIPTION_ID}/subscriptions",
+        "get",
+        response=aresponses.Response(text="Conflict", status=409),
+    )
     server.add(
         "api.simplisafe.com",
         f"/v1/users/{TEST_SUBSCRIPTION_ID}/subscriptions",
@@ -349,6 +324,13 @@ async def test_request_error_successful_retry(
             TEST_AUTHORIZATION_CODE, TEST_CODE_VERIFIER, session=session
         )
 
+        simplisafe.disable_request_retries()
+
+        with pytest.raises(RequestError):
+            await simplisafe.async_get_systems()
+
+        simplisafe.enable_request_retries()
+
         # If this succeeds without throwing an exception, the retry is successful:
         await simplisafe.async_get_systems()