Add retry mechanism for 429 Too Many Requests and 500 Internal Server Error

lib/reservation_monitor.py

@@ -20,6 +20,9 @@
 from .webdriver import WebDriver
 
 TOO_MANY_REQUESTS_CODE = 429
+INTERNAL_SERVER_ERROR_CODE = 500
+
+RETRY_WAIT_SECONDS = 20
 
 logger = get_logger(__name__)
 
@@ -194,42 +197,63 @@ def _check(self) -> bool:
         # this scope
         return False
 
-    def _get_reservations(self) -> Tuple[List[Dict[str, Any]], bool]:
+    def _get_reservations(self, max_retries: int = 1) -> Tuple[List[Dict[str, Any]], bool]:
         """
-        Returns a list of reservations and a boolean indicating if reservation
-        scheduling should be skipped.
+        Attempts to retrieve a list of reservations and returns a tuple containing the list
+        of reservations and a boolean indicating whether reservation scheduling should be skipped.
 
-        Reservation scheduling will be skipped if a Too Many Requests error or timeout occurs
-        because new headers might not be valid and a list of reservations could not be retrieved.
+        The method will retry fetching reservations once in case of a timeout
+        or a Too Many Requests error. If the retry fails, reservation scheduling will be
+        skipped until the next scheduled attempt.
         """
-        logger.debug("Retrieving reservations for account")
-        webdriver = WebDriver(self.checkin_scheduler)
+        logger.debug("Retrieving reservations for account (max retries: %d)", max_retries)
 
-        try:
-            reservations = webdriver.get_reservations(self)
-        except DriverTimeoutError:
-            logger.debug(
-                "Timeout while retrieving reservations during login. Skipping reservation retrieval"
-            )
-            self.notification_handler.timeout_during_retrieval("account")
-            return [], True
-        except LoginError as err:
-            if err.status_code == TOO_MANY_REQUESTS_CODE:
-                # Don't exit when a Too Many Requests error happens. Instead, just skip the
-                # retrieval until the next time.
+        for attempt in range(max_retries + 1):
+            webdriver = WebDriver(self.checkin_scheduler)
+
+            try:
+                reservations = webdriver.get_reservations(self)
                 logger.debug(
-                    "Encountered a Too Many Requests error while logging in. Skipping reservation "
-                    "retrieval"
+                    "Successfully retrieved %d reservations after %d attempts",
+                    len(reservations),
+                    attempt + 1,
                 )
-                self.notification_handler.too_many_requests_during_login()
-                return [], True
-
-            logger.debug("Error logging in. %s. Exiting", err)
-            self.notification_handler.failed_login(err)
-            sys.exit(1)
-
-        logger.debug("Successfully retrieved %d reservations", len(reservations))
-        return reservations, False
+                return reservations, False
+
+            except DriverTimeoutError:
+                if attempt < max_retries:
+                    logger.debug("Timeout while retrieving reservations during login. Retrying")
+                    logger.debug("Waiting for %d seconds before retrying", RETRY_WAIT_SECONDS)
+                    time.sleep(RETRY_WAIT_SECONDS)
+                else:
+                    logger.debug(
+                        "Timeout persisted after %d retries. Skipping reservation retrieval",
+                        max_retries,
+                    )
+                    self.notification_handler.timeout_during_retrieval("account")
+
+            except LoginError as err:
+                if err.status_code in [TOO_MANY_REQUESTS_CODE, INTERNAL_SERVER_ERROR_CODE]:
+                    if attempt < max_retries:
+                        logger.debug(
+                            "Encountered an error (status: %d) while logging in. Retrying",
+                            err.status_code,
+                        )
+                        logger.debug("Waiting for %d seconds before retrying", RETRY_WAIT_SECONDS)
+                        time.sleep(RETRY_WAIT_SECONDS)
+                    else:
+                        logger.debug(
+                            "Error (status: %d) persists. Skipping reservation retrieval",
+                            err.status_code,
+                        )
+                        self.notification_handler.too_many_requests_during_login()
+                else:
+                    logger.debug("Error logging in. %s. Exiting", err)
+                    self.notification_handler.failed_login(err)
+                    time.sleep(1)
+                    sys.exit(1)
+
+        return [], True
 
     def _stop_monitoring(self) -> None:
         print(f"\nStopping monitoring for account with username {self.username}")

tests/unit/test_reservation_monitor.py

@@ -254,6 +254,7 @@ def test_check_skips_scheduling_if_an_error_occurs(self, mocker: MockerFixture)
     def test_get_reservations_skips_retrieval_on_driver_timeout(
         self, mocker: MockerFixture
     ) -> None:
+        mocker.patch("time.sleep")
         mocker.patch.object(WebDriver, "get_reservations", side_effect=DriverTimeoutError)
         mock_timeout_notif = mocker.patch.object(NotificationHandler, "timeout_during_retrieval")
 
@@ -266,6 +267,7 @@ def test_get_reservations_skips_retrieval_on_driver_timeout(
     def test_get_reservations_skips_retrieval_on_too_many_requests_error(
         self, mocker: MockerFixture
     ) -> None:
+        mocker.patch("time.sleep")
         mocker.patch.object(
             WebDriver, "get_reservations", side_effect=LoginError("", TOO_MANY_REQUESTS_CODE)
         )
@@ -280,6 +282,7 @@ def test_get_reservations_skips_retrieval_on_too_many_requests_error(
         mock_too_many_requests_notif.assert_called_once()
 
     def test_get_reservations_exits_on_login_error(self, mocker: MockerFixture) -> None:
+        mocker.patch("time.sleep")
         mocker.patch.object(WebDriver, "get_reservations", side_effect=LoginError("", 400))
         mock_failed_login = mocker.patch.object(NotificationHandler, "failed_login")