From f4cdf266004fb5c14fca5d822c4aab70f7d27c89 Mon Sep 17 00:00:00 2001
From: mattsb42 <matt.s.b.42@gmail.com>
Date: Wed, 30 Oct 2019 21:27:07 -0700
Subject: [PATCH] add single-retry GitHub credentials logic

---
 src/github_proxy.py            | 49 ++++++++++++++++++++--------------
 test/unit/test_github_proxy.py | 31 +++++++++++++++++++--
 2 files changed, 58 insertions(+), 22 deletions(-)

diff --git a/src/github_proxy.py b/src/github_proxy.py
index 2af25dc..6220bee 100644
--- a/src/github_proxy.py
+++ b/src/github_proxy.py
@@ -3,7 +3,7 @@
 import re
 
 import boto3
-from github import Github
+from github import Github, GithubException
 
 import config
 import lambdalogging
@@ -33,10 +33,6 @@
 class GithubProxy:
     """Encapsulate interactions with Github."""
 
-    def __init__(self):
-        """Initialize proxy."""
-        pass
-
     def publish_pr_comment(self, build):
         """Publish PR comment with link to build logs."""
         pr_comment = PR_COMMENT_TEMPLATE.format(
@@ -46,22 +42,35 @@ def publish_pr_comment(self, build):
             logs_url=build.get_logs_url(),
         )
 
-        # initialize client before logging to ensure GitHub attributes are populated
-        gh_client = self._get_client()
-        LOG.debug('Publishing PR Comment: repo=%s/%s, pr_id=%s, comment=%s',
-                  self._github_owner, self._github_repo, build.get_pr_id(), pr_comment)
-
-        repo = gh_client.get_user(self._github_owner).get_repo(self._github_repo)
-        repo.get_pull(build.get_pr_id()).create_issue_comment(pr_comment)
-
-    def _get_client(self):
-        if not hasattr(self, '_client'):
-            self._init_client()
-        return self._client
-
-    def _init_client(self):
+        # initialize client to obtain user info if first run
+        _ = self._github_client
+        attempt = 1
+        while attempt <= 2:
+            try:
+                LOG.debug('Publishing PR Comment: repo=%s/%s, pr_id=%s, comment=%s',
+                          self._github_owner, self._github_repo, build.get_pr_id(), pr_comment)
+
+                repo = self._github_client.get_user(self._github_owner).get_repo(self._github_repo)
+                repo.get_pull(build.get_pr_id()).create_issue_comment(pr_comment)
+            except GithubException:
+                self._refresh_github_client()
+                attempt += 1
+            else:
+                return
+
+        raise PermissionError("Unable to authenticate to GitHub with available credentials!")
+
+    def _refresh_github_client(self):
         self._init_github_info()
-        self._client = Github(self._github_token)
+        self.__github_client = Github(self._github_token)
+        return self.__github_client
+
+    @property
+    def _github_client(self):
+        try:
+            return self.__github_client
+        except AttributeError:
+            return self._refresh_github_client()
 
     def _init_github_info(self):
         response = CODEBUILD.batch_get_projects(
diff --git a/test/unit/test_github_proxy.py b/test/unit/test_github_proxy.py
index 23530fe..192633c 100644
--- a/test/unit/test_github_proxy.py
+++ b/test/unit/test_github_proxy.py
@@ -1,5 +1,7 @@
 import pytest
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, Mock
+
+from github import BadCredentialsException
 
 import github_proxy
 import test_constants
@@ -61,11 +63,27 @@ def mock_github(mocker):
     return github_proxy.Github.return_value
 
 
-def test_publish_pr_comment(mocker, mock_config, mock_codebuild, mock_secretsmanager, mock_github):
+@pytest.fixture
+def mock_renewing_github(mocker):
+    mocker.patch.object(github_proxy, 'Github')
+
+    mock_failure = Mock()
+    mock_failure.get_user.side_effect = BadCredentialsException(status='foo', data='bar')
+
+    github_proxy.Github.side_effect = (mock_failure, Mock())
+    return github_proxy.Github
+
+
+def _mock_build():
     build = MagicMock(status=BUILD_STATUS)
     build.get_logs_url.return_value = LOGS_URL
     build.get_pr_id.return_value = PR_ID
     build.commit_id = COMMIT_ID
+    return build
+
+
+def test_publish_pr_comment(mocker, mock_config, mock_codebuild, mock_secretsmanager, mock_github):
+    build = _mock_build()
 
     proxy = github_proxy.GithubProxy()
     proxy.publish_pr_comment(build)
@@ -91,6 +109,15 @@ def test_publish_pr_comment(mocker, mock_config, mock_codebuild, mock_secretsman
     mock_pr.create_issue_comment.assert_called_once_with(expected_comment)
 
 
+def test_public_pr_comment_renew_credentials(mock_config, mock_codebuild, mock_secretsmanager, mock_renewing_github):
+    build = _mock_build()
+
+    proxy = github_proxy.GithubProxy()
+    proxy.publish_pr_comment(build)
+
+    assert mock_renewing_github.call_count == 2
+
+
 def test_init_github_info_auth_with_secrets_manager_arn(mocker, mock_config, mock_codebuild, mock_secretsmanager):
     secret_arn = 'arn:secret'
     mock_config.configure_mock(GITHUB_OAUTH_TOKEN_SECRET_ARN=secret_arn)