From cfdc3d4a442ec3d715771d97f66eac96c51ae0bd Mon Sep 17 00:00:00 2001
From: Jacob Riddle <87780794+jriddle-linode@users.noreply.github.com>
Date: Mon, 14 Aug 2023 16:45:28 -0400
Subject: [PATCH] new: retry on common errors (#505)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## 📝 Description

**What does this PR do and why is this change necessary?**

Adding retry functionality for common random errors sent by the API

- 400: Bad Request with nginx headers
- 408: Request Timeout
- 429: Rate limit exceeded

## ✔️ How to Test

**What are the steps to reproduce the issue or verify the changes?**

**How do I run the relevant unit/integration tests?**

`make test`

---------

Co-authored-by: Lena Garber <114949949+lgarber-akamai@users.noreply.github.com>
Co-authored-by: Ania Misiorek <139170033+amisiorek-akamai@users.noreply.github.com>
---
 README.rst                     |  8 ++++++
 linodecli/__init__.py          |  3 ++-
 linodecli/api_request.py       | 31 +++++++++++++++++++++++
 linodecli/arg_helpers.py       |  5 ++++
 tests/unit/test_api_request.py | 45 ++++++++++++++++++++++++++++++++++
 tests/unit/test_cli.py         |  8 +++++-
 6 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index f8d033bec..e3b145108 100644
--- a/README.rst
+++ b/README.rst
@@ -220,6 +220,14 @@ in addition to its normal output.  If these warnings can interfere with your
 scripts or you otherwise want them disabled, simply add the ``--suppress-warnings``
 flag to prevent them from being emitted.
 
+Suppressing Warnings
+""""""""""""""""""""
+
+Sometimes the API responds with a error that can be ignored. For example a timeout
+or nginx response that can't be parsed correctly, by default the CLI will retry
+calls on these errors we've identified. If you'd like to disable this behavior for
+any reason use the ``--no-retry`` flag.
+
 Shell Completion
 """"""""""""""""
 
diff --git a/linodecli/__init__.py b/linodecli/__init__.py
index 95055943e..e8cdf8dc0 100755
--- a/linodecli/__init__.py
+++ b/linodecli/__init__.py
@@ -24,7 +24,6 @@
 )
 from .cli import CLI
 from .completion import bake_completions, get_completions
-from .configuration import ENV_TOKEN_NAME
 from .helpers import handle_url_overrides
 from .output import OutputMode
 
@@ -83,6 +82,8 @@ def main():  # pylint: disable=too-many-branches,too-many-statements
         cli.output_handler.columns = parsed.format
 
     cli.defaults = not parsed.no_defaults
+    cli.retry_count = 0
+    cli.no_retry = parsed.no_retry
     cli.suppress_warnings = parsed.suppress_warnings
 
     if parsed.all_columns or parsed.all:
diff --git a/linodecli/api_request.py b/linodecli/api_request.py
index 11cfeed10..99bc08091 100644
--- a/linodecli/api_request.py
+++ b/linodecli/api_request.py
@@ -5,6 +5,7 @@
 import itertools
 import json
 import sys
+import time
 from sys import version_info
 from typing import Iterable, List, Optional
 
@@ -92,6 +93,11 @@ def do_request(
     if ctx.debug_request:
         _print_response_debug_info(result)
 
+    if _check_retry(result) and ctx.no_retry and ctx.retry_count < 3:
+        time.sleep(_get_retry_after(result.headers))
+        ctx.retry_count += 1
+        do_request(ctx, operation, args, filter_header, skip_error_handling)
+
     _attempt_warn_old_version(ctx, result)
 
     if not 199 < result.status_code < 399 and not skip_error_handling:
@@ -361,3 +367,28 @@ def _handle_error(ctx, response):
             columns=["field", "reason"],
         )
     sys.exit(1)
+
+
+def _check_retry(response):
+    """
+    Check for valid retry scenario, returns true if retry is valid
+    """
+    if response.status_code in (408, 429):
+        # request timed out or rate limit exceeded
+        return True
+
+    if (
+        response.headers
+        and response.status_code == 400
+        and response.headers.get("Server") == "nginx"
+        and response.headers.get("Content-Type") == "text/html"
+    ):
+        # nginx html response
+        return True
+
+    return False
+
+
+def _get_retry_after(headers):
+    retry_str = headers.get("Retry-After", "")
+    return int(retry_str) if retry_str else 0
diff --git a/linodecli/arg_helpers.py b/linodecli/arg_helpers.py
index 94214ea44..441bd1200 100644
--- a/linodecli/arg_helpers.py
+++ b/linodecli/arg_helpers.py
@@ -126,6 +126,11 @@ def register_args(parser):
         default=False,
         help="Prevent the truncation of long values in command outputs.",
     )
+    parser.add_argument(
+        "--no-retry",
+        action="store_true",
+        help="Skip retrying on common errors like timeouts.",
+    )
     parser.add_argument(
         "--column-width",
         type=int,
diff --git a/tests/unit/test_api_request.py b/tests/unit/test_api_request.py
index b979ed073..ecf8dbd7a 100644
--- a/tests/unit/test_api_request.py
+++ b/tests/unit/test_api_request.py
@@ -8,6 +8,7 @@
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
+import pytest
 import requests
 
 from linodecli import api_request
@@ -475,3 +476,47 @@ def json_func():
 
         output = stderr_buf.getvalue()
         assert "" == output
+
+    def test_do_request_recursion(self, mock_cli, list_operation):
+        mock_response = Mock(status_code=408)
+        with patch(
+            "linodecli.api_request.requests.get", return_value=mock_response
+        ) and pytest.raises(SystemExit):
+            _ = api_request.do_request(mock_cli, list_operation, None)
+            assert mock_cli.retry_count == 3
+
+    def test_check_retry(self):
+        mock_response = Mock(status_code=200)
+        output = api_request._check_retry(mock_response)
+        assert not output
+
+        mock_response = Mock(status_code=408)
+        output = api_request._check_retry(mock_response)
+        assert output
+
+        mock_response = Mock(status_code=429)
+        output = api_request._check_retry(mock_response)
+        assert output
+
+        mock_response = Mock(
+            status_code=400,
+            headers={
+                "Server": "nginx",
+                "Content-Type": "text/html",
+            },
+        )
+        output = api_request._check_retry(mock_response)
+        assert output
+
+    def test_get_retry_after(self):
+        headers = {"Retry-After": "10"}
+        output = api_request._get_retry_after(headers)
+        assert output == 10
+
+        headers = {"Retry-After": ""}
+        output = api_request._get_retry_after(headers)
+        assert output == 0
+
+        headers = {}
+        output = api_request._get_retry_after(headers)
+        assert output == 0
diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py
index 0fe4c130a..eddfbc4db 100644
--- a/tests/unit/test_cli.py
+++ b/tests/unit/test_cli.py
@@ -16,10 +16,16 @@
 
 class MockResponse:
     def __init__(
-        self, page: int, pages: int, results: int, status_code: int = 200
+        self,
+        page: int,
+        pages: int,
+        results: int,
+        status_code: int = 200,
+        headers: dict = None,
     ):
         self.page = page
         self.pages = pages
+        self.headers = headers
         self.results = results
         self.status_code = status_code