Handle GitHub App token connect timeouts safely

src/github_app.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import contextlib
+import logging
 import time
 from typing import Generator
 
@@ -12,27 +13,56 @@
 
 
 class GithubAppToken:
+    TOKEN_REQUEST_TIMEOUT_SECONDS = 10
+    TOKEN_REQUEST_MAX_ATTEMPTS = 3
+    TOKEN_REQUEST_RETRY_BACKOFF_SECONDS = 1
+
     def __init__(self, private_key, app_id) -> None:
         self.headers = self.get_authentication_header(private_key, app_id)
 
+    def _fetch_installation_access_token(self, installation_id: int):
+        for attempt in range(1, self.TOKEN_REQUEST_MAX_ATTEMPTS + 1):
+            try:
+                req = requests.post(
+                    url=f"https://api.github.com/app/installations/{installation_id}/access_tokens",
+                    headers=self.headers,
+                    timeout=self.TOKEN_REQUEST_TIMEOUT_SECONDS,
+                )
+                req.raise_for_status()
+                return req.json()
+            except requests.exceptions.ConnectTimeout:
+                if attempt >= self.TOKEN_REQUEST_MAX_ATTEMPTS:
+                    raise
+                logging.warning(
+                    "Connect timeout requesting GitHub App token for installation %s; retrying (%s/%s)",
+                    installation_id,
+                    attempt,
+                    self.TOKEN_REQUEST_MAX_ATTEMPTS,
+                )
+                # Linear backoff avoids hot-looping when GitHub API is flaky.
+                time.sleep(self.TOKEN_REQUEST_RETRY_BACKOFF_SECONDS * attempt)
+
     # From docs: Installation access tokens have the permissions
     # configured by the GitHub App and expire after one hour.
     @contextlib.contextmanager
     def get_token(self, installation_id: int) -> Generator[str, None, None]:
-        req = requests.post(
-            url=f"https://api.github.com/app/installations/{installation_id}/access_tokens",
-            headers=self.headers,
-        )
-        req.raise_for_status()
-        resp = req.json()
+        resp = self._fetch_installation_access_token(installation_id)
         try:
             # This token expires in an hour
             yield resp["token"]
         finally:
-            requests.delete(
-                "https://api.github.com/installation/token",
-                headers={"Authorization": f"token {resp['token']}"},
-            )
+            try:
+                requests.delete(
+                    "https://api.github.com/installation/token",
+                    headers={"Authorization": f"token {resp['token']}"},
+                    timeout=self.TOKEN_REQUEST_TIMEOUT_SECONDS,
+                )
+            except requests.exceptions.RequestException:
+                logging.warning(
+                    "Failed to revoke temporary GitHub App token for installation %s",
+                    installation_id,
+                    exc_info=True,
+                )
 
     def get_jwt_token(self, private_key, app_id):
         payload = {

tests/test_github_app.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from unittest.mock import Mock
+from unittest.mock import patch
+
+import pytest
+import requests
+
+from src.github_app import GithubAppToken
+
+
+@patch.object(GithubAppToken, "get_authentication_header", return_value={})
+@patch("src.github_app.requests.delete")
+@patch("src.github_app.time.sleep")
+@patch("src.github_app.requests.post")
+def test_get_token_retries_connect_timeout_then_succeeds(
+    mock_post,
+    mock_sleep,
+    mock_delete,
+    _mock_headers,
+):
+    response = Mock()
+    response.raise_for_status.return_value = None
+    response.json.return_value = {"token": "temporary-token"}
+    mock_post.side_effect = [requests.exceptions.ConnectTimeout(), response]
+
+    token_provider = GithubAppToken(private_key="irrelevant", app_id=123)
+
+    with token_provider.get_token(installation_id=42) as token:
+        assert token == "temporary-token"
+
+    assert mock_post.call_count == 2
+    first_call = mock_post.call_args_list[0]
+    second_call = mock_post.call_args_list[1]
+    assert first_call.kwargs["timeout"] == GithubAppToken.TOKEN_REQUEST_TIMEOUT_SECONDS
+    assert second_call.kwargs["timeout"] == GithubAppToken.TOKEN_REQUEST_TIMEOUT_SECONDS
+    assert first_call.kwargs["url"].endswith("/42/access_tokens")
+    assert second_call.kwargs["url"].endswith("/42/access_tokens")
+    mock_sleep.assert_called_once_with(1)
+    mock_delete.assert_called_once()
+    assert (
+        mock_delete.call_args.kwargs["timeout"]
+        == GithubAppToken.TOKEN_REQUEST_TIMEOUT_SECONDS
+    )
+
+
+@patch.object(GithubAppToken, "get_authentication_header", return_value={})
+@patch("src.github_app.requests.delete")
+@patch("src.github_app.time.sleep")
+@patch("src.github_app.requests.post")
+def test_get_token_raises_after_repeated_connect_timeouts(
+    mock_post,
+    mock_sleep,
+    mock_delete,
+    _mock_headers,
+):
+    mock_post.side_effect = [requests.exceptions.ConnectTimeout()] * 3
+    token_provider = GithubAppToken(private_key="irrelevant", app_id=123)
+
+    with pytest.raises(requests.exceptions.ConnectTimeout):
+        with token_provider.get_token(installation_id=42):
+            pass
+
+    assert mock_post.call_count == GithubAppToken.TOKEN_REQUEST_MAX_ATTEMPTS
+    mock_sleep.assert_any_call(1)
+    mock_sleep.assert_any_call(2)
+    assert mock_sleep.call_count == 2
+    mock_delete.assert_not_called()
+
+
+@patch.object(GithubAppToken, "get_authentication_header", return_value={})
+@patch("src.github_app.requests.delete")
+@patch("src.github_app.requests.post")
+def test_get_token_ignores_delete_errors(
+    mock_post,
+    mock_delete,
+    _mock_headers,
+):
+    response = Mock()
+    response.raise_for_status.return_value = None
+    response.json.return_value = {"token": "temporary-token"}
+    mock_post.return_value = response
+    mock_delete.side_effect = requests.exceptions.Timeout()
+
+    token_provider = GithubAppToken(private_key="irrelevant", app_id=123)
+
+    with token_provider.get_token(installation_id=42) as token:
+        assert token == "temporary-token"
+
+    mock_delete.assert_called_once()