[MergeQueues] Delete branch if PR is merged (#193)

### Summary

Github's delete branch logic doesn't work if SGTM added the pull request to the merge queue. This will make SGTM delete branches if they haven't been cleaned up yet.

Asana tasks:
https://app.asana.com/0/1125126232217429/1207661000934486/f


Relevant deployment: 

CC: @suzyng83209 @prebeta @vn6 @michael-huang87

### Test Plan

Run on test sgtm deployment

### Risks



Pull Request: #193



Pull Request synchronized with [Asana task](https://app.asana.com/0/0/1207680708784404)

Author: prebeta

src/github/client.py

@@ -1,40 +1,41 @@
 import requests
 from requests.auth import HTTPBasicAuth
 
-from github import PullRequest  # type: ignore
+from github.PullRequest import PullRequest
+from github.GithubException import GithubException
 from src.github.get_app_token import sgtm_github_auth
 from src.logger import logger
 
 gh_client = sgtm_github_auth.get_rest_client()
 
 
-def _get_pull_request(owner: str, repository: str, number: int) -> PullRequest:  # type: ignore
+def _get_pull_request(owner: str, repository: str, number: int) -> PullRequest:
     repo = gh_client.get_repo(f"{owner}/{repository}")
     pr = repo.get_pull(number)
-    return pr  # type: ignore
+    return pr
 
 
 def edit_pr_description(owner: str, repository: str, number: int, description: str):
     pr = _get_pull_request(owner, repository, number)
-    pr.edit(body=description)  # type: ignore
+    pr.edit(body=description)
 
 
 def edit_pr_title(owner: str, repository: str, number: int, title: str):
     pr = _get_pull_request(owner, repository, number)
-    pr.edit(title=title)  # type: ignore
+    pr.edit(title=title)
 
 
 def add_pr_comment(owner: str, repository: str, number: int, comment: str):
     pr = _get_pull_request(owner, repository, number)
-    pr.create_issue_comment(comment)  # type: ignore
+    pr.create_issue_comment(comment)
 
 
 def set_pull_request_assignee(owner: str, repository: str, number: int, assignee: str):
     repo = gh_client.get_repo(f"{owner}/{repository}")
     # Using get_issue here because get_pull returns a pull request which only
     # allows you to *add* an assignee, not set the assignee.
     pr = repo.get_issue(number)
-    pr.edit(assignee=assignee)  # type: ignore
+    pr.edit(assignee=assignee)
 
 
 def merge_pull_request(owner: str, repository: str, number: int, title: str, body: str):
@@ -44,13 +45,15 @@ def merge_pull_request(owner: str, repository: str, number: int, title: str, bod
     # which we rely on for code review tests.
     title_with_number = f"{title} (#{number})"
     try:
-        pr.enable_automerge(commit_headline=title_with_number, commit_body=body)  # type: ignore
+        pr.enable_automerge(commit_headline=title_with_number, commit_body=body)
     except Exception as e:
         logger.info(
             f"Failed to enable automerge for PR {title_with_number}, with error {e}"
         )
         logger.info("Merging PR manually")
-        pr.merge(commit_title=title_with_number, commit_message=body, merge_method="squash")  # type: ignore
+        pr.merge(
+            commit_title=title_with_number, commit_message=body, merge_method="squash"
+        )
 
 
 def rerequest_check_run(owner: str, repository: str, check_run_id: int):
@@ -60,3 +63,28 @@ def rerequest_check_run(owner: str, repository: str, check_run_id: int):
     )
     # Some check runs cannot be rerequested. See https://docs.github.com/en/rest/checks/runs?apiVersion=2022-11-28#rerequest-a-check-run--status-codes
     return requests.post(url, auth=auth).status_code == 201
+
+
+def delete_branch_if_exists(owner: str, repo_name: str, branch_name: str):
+    """
+    Deletes a branch from a GitHub repository if it exists.
+
+    Args:
+        owner (str): The owner of the repository.
+        repo_name (str): The name of the repository.
+        branch_name (str): The name of the branch to delete.
+    """
+    try:
+        repo = gh_client.get_repo(f"{owner}/{repo_name}")
+        # Attempt to get the branch, will raise a 404 error if not found
+        repo.get_branch(branch_name)
+        ref = f"heads/{branch_name}"
+        git_ref = repo.get_git_ref(ref)
+        git_ref.delete()
+        logger.info(f"Branch '{branch_name}' deleted successfully.")
+    except GithubException as e:
+        if e.status == 404:
+            logger.info(f"Branch '{branch_name}' does not exist or is already deleted.")
+        else:
+            logger.error(f"Error deleting branch: {e}")
+            raise

src/github/graphql/fragments/FullPullRequest.py

@@ -6,6 +6,7 @@
 fragment FullPullRequest on PullRequest {
   id
   baseRefName
+  headRefName
   body
   bodyHTML
   title

src/github/logic.py

@@ -296,6 +296,14 @@ def maybe_automerge_pull_request(pull_request: PullRequest) -> bool:
     return False
 
 
+def maybe_delete_branch_if_merged(pull_request: PullRequest):
+    if pull_request.merged():
+        owner = pull_request.repository_owner_handle()
+        repo_name = pull_request.repository_name()
+        branch_name = pull_request.head_ref_name()
+        github_client.delete_branch_if_exists(owner, repo_name, branch_name)
+
+
 # ----------------------------------------------------------------------------------
 # Automerge helpers
 # ----------------------------------------------------------------------------------

src/github/models/pull_request.py

@@ -219,3 +219,6 @@ def labels(self) -> List[Label]:
 
     def base_ref_name(self) -> str:
         return self._raw["baseRefName"]
+
+    def head_ref_name(self) -> str:
+        return self._raw["headRefName"]

src/github/webhook.py

@@ -18,6 +18,8 @@ def _handle_pull_request_webhook(payload: dict) -> HttpResponse:
         pull_request = graphql_client.get_pull_request(pull_request_id)
         # a label change will trigger this webhook, so it may trigger automerge
         github_logic.maybe_automerge_pull_request(pull_request)
+        if payload["action"] == "closed":
+            github_logic.maybe_delete_branch_if_merged(pull_request)
         github_logic.maybe_add_automerge_warning_comment(pull_request)
         github_controller.upsert_pull_request(pull_request)
         return HttpResponse("200")

test/github/test_logic.py

@@ -602,6 +602,27 @@ def test_does_not_add_warning_comment_if_pr_is_approved(
         add_pr_comment_mock.assert_not_called()
 
 
+@patch.object(github_client, "delete_branch_if_exists")
+class TestMaybeDeleteBranchIfMerged(unittest.TestCase):
+    def test_maybe_delete_branch_if_merged(self, mock_delete_branch):
+        pull_request = build(
+            builder.pull_request().merged(True).head_ref_name("feature-branch")
+        )
+        github_logic.maybe_delete_branch_if_merged(pull_request)
+        mock_delete_branch.assert_called_once_with(
+            pull_request.repository_owner_handle(),
+            pull_request.repository_name(),
+            "feature-branch",
+        )
+
+    def test_do_not_delete_branch_if_closed(self, mock_delete_branch):
+        pull_request = build(
+            builder.pull_request().merged(False).head_ref_name("feature-branch")
+        )
+        github_logic.maybe_delete_branch_if_merged(pull_request)
+        mock_delete_branch.assert_not_called()
+
+
 if __name__ == "__main__":
     from unittest import main as run_tests
 

test/github/test_webhook.py

@@ -123,6 +123,38 @@ def test_comment_deletion_when_review_not_found(
         delete_comment.assert_called_once_with(self.COMMENT_NODE_ID)
 
 
+@patch("src.github.controller.upsert_pull_request")
+@patch("src.github.logic.maybe_automerge_pull_request")
+@patch("src.github.logic.maybe_delete_branch_if_merged")
+@patch("src.github.graphql.client.get_pull_request")
+class TestHandlePullRequestWebhookClosed(MockDynamoDbTestCase):
+    PULL_REQUEST_NODE_ID = "abcdef"
+
+    def test_handle_pull_request_webhook_when_closed(
+        self,
+        get_pull_request,
+        maybe_delete_branch_if_merged,
+        maybe_automerge_pull_request,
+        upsert_pull_request,
+    ):
+        payload = {
+            "action": "closed",
+            "pull_request": {
+                "node_id": self.PULL_REQUEST_NODE_ID,
+            },
+        }
+
+        pull_request = MagicMock(spec=PullRequest)
+        get_pull_request.return_value = pull_request
+
+        response = webhook._handle_pull_request_webhook(payload)
+        self.assertEqual(response.status_code, "200")
+
+        maybe_automerge_pull_request.assert_called_once()
+        maybe_delete_branch_if_merged.assert_called_once()
+        upsert_pull_request.assert_called_once()
+
+
 if __name__ == "__main__":
     from unittest import main as run_tests
 

test/impl/builders/pull_request_builder.py

@@ -27,6 +27,7 @@ def __init__(self, body: str = ""):
             "number": pr_number,
             "body": body,
             "baseRefName": create_uuid(),
+            "headRefName": create_uuid(),
             "title": create_uuid(),
             "url": "https://www.github.com/foo/pulls/" + str(pr_number),
             "assignees": {"nodes": []},
@@ -168,6 +169,10 @@ def base_ref_name(self, base_ref_name: str):
         self.raw_pr["baseRefName"] = base_ref_name
         return self
 
+    def head_ref_name(self, head_ref_name: str):
+        self.raw_pr["headRefName"] = head_ref_name
+        return self
+
     def build(self) -> PullRequest:
         return PullRequest(self.raw_pr)