fix: treat 401 re-authenticate correctly

GuyLewin · Feb 26, 2022 · 496fc64 · 496fc64
1 parent 30b6414
commit 496fc64
Showing 1 changed file with 19 additions and 11 deletions.
diff --git a/custom_components/lifetime_fitness/api.py b/custom_components/lifetime_fitness/api.py
@@ -1,6 +1,8 @@
 """API client for Life Time Fitness"""
-from aiohttp import ClientError
+from aiohttp import ClientError, ClientResponseError
 from datetime import date
+from http import HTTPStatus
+import logging
 
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 
@@ -18,6 +20,8 @@
     API_CLUB_VISITS_AUTH_HEADER,
 )
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class Api:
     def __init__(self, hass, username: str, password: str) -> None:
@@ -55,17 +59,20 @@ async def _get_visits_between_dates(self, start_date: date, end_date: date):
         if self._sso_token is None:
             raise ApiAuthRequired
 
-        async with self._clientsession.get(
-            API_CLUB_VISITS_ENDPOINT_FORMATSTRING.format(
-                start_date=start_date.strftime(API_CLUB_VISITS_ENDPOINT_DATE_FORMAT),
-                end_date=end_date.strftime(API_CLUB_VISITS_ENDPOINT_DATE_FORMAT),
-            ),
-            headers={API_CLUB_VISITS_AUTH_HEADER: self._sso_token},
-        ) as response:
-            if response.status == 401:
+        try:
+            async with self._clientsession.get(
+                API_CLUB_VISITS_ENDPOINT_FORMATSTRING.format(
+                    start_date=start_date.strftime(API_CLUB_VISITS_ENDPOINT_DATE_FORMAT),
+                    end_date=end_date.strftime(API_CLUB_VISITS_ENDPOINT_DATE_FORMAT),
+                ),
+                headers={API_CLUB_VISITS_AUTH_HEADER: self._sso_token},
+            ) as response:
+                response_json = await response.json()
+                return response_json
+        except ClientResponseError as err:
+            if err.status == HTTPStatus.UNAUTHORIZED:
                 raise ApiAuthExpired
-            response_json = await response.json()
-            return response_json
+            raise err
 
     async def update(self):
         today = date.today()
@@ -79,6 +86,7 @@ async def update(self):
                 self.result_json = await self._get_visits_between_dates(first_day_of_the_year, today)
         except Exception as e:
             self.update_successful = False
+            _LOGGER.exception("Unexpected exception during Life Time API update")
             raise e