Add option to checkout newly created branches

Additionally, fix an inconsistency where the create() function would
leave the branch checked out when performing a hard reset, but not for
newly created branches or for existing branches that didn't get reset.

Author: jpichon

git_wrapper/branch.py

@@ -49,7 +49,7 @@ def _run_cherry(self, upstream, head, regex):
         return ret_data
 
     @reference_exists("start_ref")
-    def create(self, name, start_ref, reset_if_exists=False):
+    def create(self, name, start_ref, reset_if_exists=False, checkout=False):
         """Create a local branch based on start_ref.
 
            If the branch exists, do nothing or hard reset it if reset_if_exists
@@ -60,14 +60,20 @@ def create(self, name, start_ref, reset_if_exists=False):
                                  a starting point.
            :param bool reset_if_exists: Whether to hard reset the branch to
                                         start_ref if the branch already exists.
+           :param bool checkout: Whether to checkout the new branch
+           :return: True if a new branch was created, None otherwise
         """
         if not self.exists(name):
             self.git_repo.git.branch(name, start_ref)
+            if checkout:
+                self.git_repo.git.checkout(name)
             return True
-        if self.exists(name) and not reset_if_exists:
-            return
+
         if self.exists(name) and reset_if_exists:
-            self.hard_reset_to_ref(name, start_ref)
+            self.hard_reset_to_ref(name, start_ref, checkout)
+
+        if checkout:
+            self.git_repo.git.checkout(name)
 
     def exists(self, name, remote=None):
         """Checks if a branch exists locally or on the specified remote.
@@ -313,11 +319,12 @@ def hard_reset(self, branch="master", remote="origin",
         remote_ref = "{0}/{1}".format(remote, remote_branch)
         self.hard_reset_to_ref(branch, remote_ref)
 
-    def hard_reset_to_ref(self, branch, ref):
+    def hard_reset_to_ref(self, branch, ref, checkout=True):
         """Perform a hard reset of a local branch to any reference.
 
            :param str branch: Local branch to reset
            :param str ref: Reference (commit, tag, ...) to reset to
+           :param bool checkout: Whether to checkout the new branch
         """
         # Ensure the reference maps to a commit
         try:
@@ -326,6 +333,13 @@ def hard_reset_to_ref(self, branch, ref):
             msg = "Could not find reference {0}.".format(ref)
             raise_from(exceptions.ReferenceNotFoundException(msg), ex)
 
+        try:
+            # Preserve the reference name if there is one
+            orig = self.git_repo.repo.head.ref.name
+        except TypeError:
+            # Detached head
+            orig = self.git_repo.repo.head.commit.hexsha
+
         # Switch to the branch
         try:
             self.git_repo.git.checkout(branch)
@@ -348,6 +362,17 @@ def hard_reset_to_ref(self, branch, ref):
             )
             raise_from(exceptions.ResetException(msg), ex)
 
+        # Return to the original head if required
+        if not checkout:
+            try:
+                self.git_repo.git.checkout(orig)
+            except git.GitCommandError as ex:
+                msg = (
+                    "Could not checkout {orig}. Error: {error}".format(
+                        orig=orig, error=ex)
+                )
+                raise_from(exceptions.CheckoutException(msg), ex)
+
     @reference_exists('remote_branch')
     @reference_exists('hash_')
     def remote_contains(self, remote_branch, hash_):

integration_tests/test_branch.py

@@ -149,6 +149,28 @@ def test_create_branch(repo_root):
     assert repo.repo.branches[branch_name].commit.hexsha == tag_0_1_0_hexsha
 
 
+def test_create_and_checkout_branch(repo_root):
+    repo = GitRepo(repo_root)
+    branch_name = "test_create"
+
+    assert repo.repo.active_branch.name == 'master'
+
+    # Create and check out the new branch
+    repo.branch.create(branch_name, "0.0.1", checkout=True)
+    assert repo.repo.active_branch.name == branch_name
+
+    repo.repo.heads.master.checkout()
+    assert repo.repo.active_branch.name == 'master'
+
+    # Branch already exists - reset it and don't check it out
+    repo.branch.create(branch_name, "0.1.0", True, checkout=False)
+    assert repo.repo.active_branch.name == 'master'
+
+    # Branch already exists - reset it and check it out
+    repo.branch.create(branch_name, "0.0.1", True, checkout=True)
+    assert repo.repo.active_branch.name == branch_name
+
+
 def test_remote_contains(repo_root, patch_cleanup, datadir):
     repo = GitRepo(repo_root)
     remote_branch = "origin/master"

tests/test_branch.py

@@ -680,6 +680,77 @@ def test_reset_reset_failure(mock_repo):
             repo.branch.hard_reset(refresh=False)
 
 
+def test_reset_to_ref_with_checkout(mock_repo):
+    """
+    GIVEN GitRepo is initialized with a path and repo
+    WHEN reset is called with checkout
+    THEN repo.head.reset is called
+    AND repo.checkout is called once
+    """
+    repo = GitRepo(repo=mock_repo)
+    with patch('git.repo.fun.name_to_object'):
+        repo.branch.hard_reset_to_ref("main", "origin/main", checkout=True)
+
+    assert mock_repo.head.reset.called is True
+    assert mock_repo.git.checkout.call_count == 1
+
+
+def test_reset_to_ref_detached_head_with_checkout(mock_repo, monkeypatch):
+    """
+    GIVEN GitRepo is initialized with a path and repo
+    WHEN reset is called with checkout
+    AND the current HEAD is detached
+    THEN repo.head.reset is called
+    AND repo.checkout is called once
+    """
+    class MockRef:
+        @property
+        def name(self):
+            # Detached heads don't have a name
+            raise TypeError
+
+    repo = GitRepo(repo=mock_repo)
+    with patch('git.repo.fun.name_to_object'):
+        mock_repo.head.ref = MockRef()
+        repo.branch.hard_reset_to_ref("main", "origin/main", checkout=True)
+
+    assert mock_repo.head.reset.called is True
+    assert mock_repo.git.checkout.call_count == 1
+
+
+def test_reset_to_ref_without_checkout(mock_repo):
+    """
+    GIVEN GitRepo is initialized with a path and repo
+    WHEN reset_to_ref is called with checkout False
+    THEN repo.head.reset is called
+    AND repo.checkout is called twice to return to the original state
+    """
+    repo = GitRepo(repo=mock_repo)
+    with patch('git.repo.fun.name_to_object'):
+        repo.branch.hard_reset_to_ref("main", "origin/main", checkout=False)
+
+    assert mock_repo.head.reset.called is True
+    assert mock_repo.git.checkout.call_count == 2
+
+
+def test_reset_to_ref_without_checkout_fails(mock_repo):
+    """
+    GIVEN GitRepo is initialized with a path and repo
+    WHEN reset_to_ref is called with checkout False
+    AND switching back fails
+    THEN checkoutException is raised
+    """
+    mock_repo.git.checkout.side_effect = [None, git.GitCommandError('checkout', '')]
+
+    repo = GitRepo(repo=mock_repo)
+    with patch('git.repo.fun.name_to_object'):
+        with pytest.raises(exceptions.CheckoutException):
+            repo.branch.hard_reset_to_ref("main", "origin/main", checkout=False)
+
+    assert mock_repo.head.reset.called is True
+    assert mock_repo.git.checkout.call_count == 2
+
+
 def test_local_branch_exists(mock_repo):
     """
     GIVEN GitRepo is initialized with a path and repo
@@ -751,12 +822,29 @@ def test_create_branch(mock_repo):
     GIVEN GitRepo is initialized with a path and repo
     WHEN branch.create is called with a valid name and start_ref
     THEN git.branch is called
+    AND git.checkout is not called
     """
     repo = GitRepo(repo=mock_repo)
 
     with patch('git.repo.fun.name_to_object'):
         assert repo.branch.create("test", "123456") is True
     repo.git.branch.assert_called_with("test", "123456")
+    repo.git.checkout.assert_not_called()
+
+
+def test_create_and_checkout_branch(mock_repo):
+    """
+    GIVEN GitRepo is initialized with a path and repo
+    WHEN branch.create is called with valid parameters and checkout is True
+    THEN git.branch is called
+    AND git.checkout is called
+    """
+    repo = GitRepo(repo=mock_repo)
+
+    with patch('git.repo.fun.name_to_object'):
+        assert repo.branch.create("test", "123456", checkout=True) is True
+    repo.git.branch.assert_called_with("test", "123456")
+    repo.git.checkout.assert_called()
 
 
 def test_create_branch_with_bad_start_ref(mock_repo):
@@ -786,6 +874,24 @@ def test_create_branch_already_exists(mock_repo):
     with patch('git.repo.fun.name_to_object'):
         repo.branch.create("test", "123456")
     assert repo.git.branch.called is False
+    assert repo.git.checkout.called is False
+
+
+def test_create_branch_already_exists_and_check_it_out(mock_repo):
+    """
+    GIVEN GitRepo is initialized with a path and repo
+    WHEN branch.create is called with valid params and checkout is True
+    AND the branch already exists
+    THEN git.branch is not called
+    AND git.checkout is called
+    """
+    repo = GitRepo(repo=mock_repo)
+    mock_repo.branches = ["test", "master"]
+
+    with patch('git.repo.fun.name_to_object'):
+        repo.branch.create("test", "123456", checkout=True)
+    assert repo.git.branch.called is False
+    assert repo.git.checkout.called is True
 
 
 def test_create_branch_already_exists_and_reset_it(mock_repo):