From 307b2cd4a299bac8121d268840fe7d4c0ea000e9 Mon Sep 17 00:00:00 2001
From: Grzegorz Dlugoszewski <g.dlugoszewski@gmail.com>
Date: Tue, 19 May 2020 17:32:04 +0200
Subject: [PATCH] Add branch checkout helper and test branches being ahead of
 upstream

---
 branch_test.go  | 21 ++++++++++++++++++++-
 helpers_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/branch_test.go b/branch_test.go
index 75a3424..ce27b5e 100644
--- a/branch_test.go
+++ b/branch_test.go
@@ -37,10 +37,17 @@ func TestClonedBranches(t *testing.T) {
 	stageFile(t, origin, "file")
 	createCommit(t, origin, "Initial commit")
 
+	createBranch(t, origin, "branch")
+
 	repo, err := CloneRepo(origin.Path(), newTempDir(t))
 	checkFatal(t, err)
 
-	createBranch(t, repo, "branch")
+	createBranch(t, repo, "local")
+
+	checkoutBranch(t, repo, "branch")
+	createFile(t, repo, "anotherFile")
+	stageFile(t, repo, "anotherFile")
+	createCommit(t, repo, "Second commit")
 
 	branches, err := Branches(repo)
 	checkFatal(t, err)
@@ -62,6 +69,18 @@ func TestClonedBranches(t *testing.T) {
 		{branches["branch"], BranchStatus{
 			Name:        "branch",
 			IsRemote:    false,
+			HasUpstream: true,
+			Ahead:       1,
+			NeedsPush:   true,
+		}},
+		{branches["origin/branch"], BranchStatus{
+			Name:        "origin/branch",
+			IsRemote:    true,
+			HasUpstream: false,
+		}},
+		{branches["local"], BranchStatus{
+			Name:        "local",
+			IsRemote:    false,
 			HasUpstream: false,
 		}},
 	}
diff --git a/helpers_test.go b/helpers_test.go
index 1ef375e..9234163 100644
--- a/helpers_test.go
+++ b/helpers_test.go
@@ -107,3 +107,34 @@ func createBranch(t *testing.T, repo *git.Repository, name string) *git.Branch {
 
 	return branch
 }
+
+func checkoutBranch(t *testing.T, repo *git.Repository, name string) {
+	branch, err := repo.LookupBranch(name, git.BranchAll)
+
+	// If branch can't be found, let's check if it's a remote branch
+	if branch == nil {
+		branch, err = repo.LookupBranch("origin/"+name, git.BranchAll)
+	}
+	checkFatal(t, errors.Wrap(err, "Failed looking up branch"))
+
+	// If branch is remote, we need to create a local one first
+	if branch.IsRemote() {
+		commit, err := repo.LookupCommit(branch.Target())
+		checkFatal(t, errors.Wrap(err, "Failed looking up commit"))
+
+		localBranch, err := repo.CreateBranch(name, commit, false)
+		checkFatal(t, errors.Wrap(err, "Failed creating local branch"))
+
+		err = localBranch.SetUpstream("origin/" + name)
+		checkFatal(t, errors.Wrap(err, "Failed setting upstream"))
+	}
+
+	err = repo.SetHead("refs/heads/" + name)
+	checkFatal(t, errors.Wrap(err, "Failed setting head"))
+
+	options := &git.CheckoutOpts{
+		Strategy: git.CheckoutForce,
+	}
+	err = repo.CheckoutHead(options)
+	checkFatal(t, errors.Wrap(err, "Failed checking out tree"))
+}