From 07377285b68bc20a48c5d2d12f5e90c1e14bd5db Mon Sep 17 00:00:00 2001
From: Laurent Mazuel <lmazuel@microsoft.com>
Date: Thu, 2 Aug 2018 10:52:08 -0700
Subject: [PATCH] Delete branch if PR closed or merged (#65)

* Delete branch if PR closed or merged

* Simplify regenerate with the ref API

* Test clean of regenerate
---
 swaggertosdk/restapi/github_handler.py | 30 +++++-----------
 tests/test_github_handler.py           | 47 ++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 22 deletions(-)
 create mode 100644 tests/test_github_handler.py

diff --git a/swaggertosdk/restapi/github_handler.py b/swaggertosdk/restapi/github_handler.py
index f69dc6f..6782339 100644
--- a/swaggertosdk/restapi/github_handler.py
+++ b/swaggertosdk/restapi/github_handler.py
@@ -178,6 +178,7 @@ def rest_pr_management(rest_pr, sdk_repo, sdk_tag, sdk_default_base=_DEFAULT_SDK
     sdk_pr_as_issue = sdk_repo.get_issue(sdk_pr.number)
     sdk_pr_merged = False
     if rest_pr.closed_at:  # If there is a date, this is closed
+        head_ref = sdk_repo.get_git_ref("heads/{}".format(sdk_pr_head))
         if rest_pr.merged:
             manage_labels(sdk_pr_as_issue,
                           to_add=[SwaggerToSdkLabels.merged],
@@ -187,6 +188,8 @@ def rest_pr_management(rest_pr, sdk_repo, sdk_tag, sdk_default_base=_DEFAULT_SDK
                     # Merge "single context PRs" automatically
                     sdk_pr.merge(merge_method="squash")
                     sdk_pr_merged = True
+                    # Delete branch from merged PR
+                    head_ref.delete()
                 except Exception as err:
                     _LOGGER.warning("Was unable to merge: %s", err)
         else:
@@ -194,6 +197,8 @@ def rest_pr_management(rest_pr, sdk_repo, sdk_tag, sdk_default_base=_DEFAULT_SDK
                           to_add=[SwaggerToSdkLabels.refused],
                           to_remove=[SwaggerToSdkLabels.in_progress])
             sdk_pr.edit(state="closed")
+            # Delete branch from closed PR
+            head_ref.delete()
     else:
         # Try to remove "refused", if it was re-opened
         manage_labels(sdk_pr_as_issue,
@@ -265,26 +270,7 @@ def clean_sdk_pr(rest_pr, sdk_repo):
         return "Didn't find the SDK PR"
 
     #
-    # Delete the branch. I need to clone the
+    # Delete the branch.
     #
-    gh_token = os.environ["GH_TOKEN"]
-    login = user_from_token(gh_token).login
-    credentials_part = '{user}:{token}@'.format(
-        user=login,
-        token=gh_token
-    )
-    upstream_url = 'https://{credentials}github.com/{sdk_git_id}.git'.format(
-        credentials=credentials_part,
-        sdk_git_id=sdk_pr.head.repo.full_name
-    )
-    with tempfile.TemporaryDirectory() as temp_dir, \
-            manage_git_folder(gh_token, Path(temp_dir) / Path("sdk"), sdk_repo.full_name) as sdk_folder:
-
-        sdk_repo = Repo(str(sdk_folder))
-        configure_user(gh_token, sdk_repo)
-
-        upstream = sdk_repo.create_remote('upstream', url=upstream_url)
-        upstream.fetch()
-        msg = upstream.push(":"+sdk_pr_head)  # old style delete
-
-        _LOGGER.debug(msg)
+    head_ref = sdk_repo.get_git_ref("heads/{}".format(sdk_pr_head))
+    head_ref.delete()
diff --git a/tests/test_github_handler.py b/tests/test_github_handler.py
new file mode 100644
index 0000000..8bcd42d
--- /dev/null
+++ b/tests/test_github_handler.py
@@ -0,0 +1,47 @@
+from types import SimpleNamespace
+
+import pytest
+
+from swaggertosdk.restapi.github_handler import clean_sdk_pr
+
+def test_clean_sdk_pr(github_client):
+
+    # Mock a Rest PR from a fork
+    rest_pr = SimpleNamespace(
+        number=666,
+        head=SimpleNamespace(
+            repo=None,  # Deleted fork
+        ),
+        base=SimpleNamespace(
+            repo=None  # Don't need the base repo if fork
+        ),
+    )
+
+    sdk_repo = github_client.get_repo("lmazuel/TestingRepo")
+
+    # Create a copy of branch "test_clean_base"
+    # If this branch does not exist, test will fail
+    test_clean_base_ref = sdk_repo.get_git_ref("heads/test_clean_base")
+    sdk_repo.create_git_ref(
+        "refs/heads/restapi_auto_666",
+        test_clean_base_ref.object.sha
+    )
+
+    # Create PR
+    sdk_repo.create_pull(
+        title="Testing clean",
+        body="Testing clean",
+        head="restapi_auto_666",
+        base="master"
+    )
+
+    # Actual test
+    result = clean_sdk_pr(rest_pr, sdk_repo)
+    assert result is None
+
+    # Assert branch is gone, means the PR is gone as well
+    try:
+        sdk_repo.get_git_ref("heads/restapi_auto_666")
+        pytest.fail("Should have fail, because the branch should be gone")
+    except Exception:
+        pass
\ No newline at end of file