Add ahead/behind detection, clean up tests

new/helpers_test.go

@@ -40,26 +40,101 @@ func NewRepoEmpty(t *testing.T) *TestRepo {
 }
 
 func NewRepoWithUntracked(t *testing.T) *TestRepo {
-	repo := NewRepoEmpty(t)
-	repo.NewFile("README", "I'm a README file")
+	tr := NewRepoEmpty(t)
+	tr.WriteFile("README", "I'm a README file")
 
-	return repo
+	return tr
 }
 
 func NewRepoWithStaged(t *testing.T) *TestRepo {
-	repo := NewRepoEmpty(t)
-	repo.NewFile("README", "I'm a README file")
-	repo.AddFile("README")
+	tr := NewRepoEmpty(t)
+	tr.WriteFile("README", "I'm a README file")
+	tr.AddFile("README")
 
-	return repo
+	return tr
 }
 func NewRepoWithCommit(t *testing.T) *TestRepo {
-	repo := NewRepoEmpty(t)
-	repo.NewFile("README", "I'm a README file")
-	repo.AddFile("README")
-	repo.NewCommit("Initial commit")
+	tr := NewRepoEmpty(t)
+	tr.WriteFile("README", "I'm a README file")
+	tr.AddFile("README")
+	tr.NewCommit("Initial commit")
 
-	return repo
+	return tr
+}
+
+func NewRepoWithModified(t *testing.T) *TestRepo {
+	tr := NewRepoEmpty(t)
+	tr.WriteFile("README", "I'm a README file")
+	tr.AddFile("README")
+	tr.NewCommit("Initial commit")
+	tr.WriteFile("README", "I'm modified")
+
+	return tr
+}
+
+func NewRepoWithIgnored(t *testing.T) *TestRepo {
+	tr := NewRepoEmpty(t)
+	tr.WriteFile(".gitignore", "ignoreme")
+	tr.AddFile(".gitignore")
+	tr.NewCommit("Initial commit")
+	tr.WriteFile("ignoreme", "I'm being ignored")
+
+	return tr
+}
+
+func NewRepoWithLocalBranch(t *testing.T) *TestRepo {
+	tr := NewRepoWithCommit(t)
+	tr.NewBranch("local")
+	return tr
+}
+
+func NewRepoWithClonedBranch(t *testing.T) *TestRepo {
+	origin := NewRepoWithCommit(t)
+
+	tr := origin.Clone()
+	tr.NewBranch("local")
+
+	return tr
+}
+
+func NewRepoWithBranchAhead(t *testing.T) *TestRepo {
+	origin := NewRepoWithCommit(t)
+
+	tr := origin.Clone()
+	tr.WriteFile("new", "I'm a new file")
+	tr.AddFile("new")
+	tr.NewCommit("New commit")
+
+	return tr
+}
+
+func NewRepoWithBranchBehind(t *testing.T) *TestRepo {
+	origin := NewRepoWithCommit(t)
+
+	tr := origin.Clone()
+
+	origin.WriteFile("origin.new", "I'm a new file on origin")
+	origin.AddFile("origin.new")
+	origin.NewCommit("New origin commit")
+
+	tr.Fetch()
+	return tr
+}
+
+func NewRepoWithBranchAheadAndBehind(t *testing.T) *TestRepo {
+	origin := NewRepoWithCommit(t)
+
+	tr := origin.Clone()
+	tr.WriteFile("local.new", "I'm a new file on local")
+	tr.AddFile("local.new")
+	tr.NewCommit("New local commit")
+
+	origin.WriteFile("origin.new", "I'm a new file on origin")
+	origin.AddFile("origin.new")
+	origin.NewCommit("New origin commit")
+
+	tr.Fetch()
+	return tr
 }
 
 func NewTempDir(t *testing.T) string {
@@ -77,12 +152,12 @@ func NewTempDir(t *testing.T) string {
 	return dir
 }
 
-func (r *TestRepo) NewFile(name string, content string) {
+func (r *TestRepo) WriteFile(name string, content string) {
 	wt, err := r.Repo.Worktree()
 	checkFatal(r.t, errors.Wrap(err, "Failed getting worktree"))
 
-	file, err := wt.Filesystem.Create(name)
-	checkFatal(r.t, errors.Wrap(err, "Failed creating a file"))
+	file, err := wt.Filesystem.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
+	checkFatal(r.t, errors.Wrap(err, "Failed opening a file"))
 
 	_, err = file.Write([]byte(content))
 	checkFatal(r.t, errors.Wrap(err, "Failed writing a file"))
@@ -139,36 +214,14 @@ func (r *TestRepo) Clone() *TestRepo {
 	}
 }
 
-//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"))
-//}
+func (r *TestRepo) Fetch() {
+	repo := &Repo{
+		repo: r.Repo,
+	}
+
+	err := repo.Fetch()
+	checkFatal(r.t, err)
+}
 
 func checkFatal(t *testing.T, err error) {
 	if err != nil {

new/repo.go

@@ -53,11 +53,24 @@ func OpenRepo(path string) (r *Repo, err error) {
 
 func newRepo(repo *git.Repository) *Repo {
 	return &Repo{
-		repo: repo,
-		Status: &RepoStatus{
-			HasUntrackedFiles:     false,
-			HasUncommittedChanges: false,
-			Branches:              make(map[string]*BranchStatus),
-		},
+		repo:   repo,
+		Status: &RepoStatus{},
 	}
 }
+
+// Fetch performs a git fetch on all remotes
+func (r *Repo) Fetch() error {
+	remotes, err := r.repo.Remotes()
+	if err != nil {
+		return errors.Wrap(err, "Failed getting remotes")
+	}
+
+	for _, remote := range remotes {
+		err = remote.Fetch(&git.FetchOptions{})
+		if err != nil {
+			return errors.Wrapf(err, "Failed fetching remote %s", remote.Config().Name)
+		}
+	}
+
+	return nil
+}

new/repo_test.go

@@ -1,95 +0,0 @@
-package new
-
-import (
-	"testing"
-)
-
-func TestRepoClone(t *testing.T) {
-	origin := NewRepoWithCommit(t)
-
-	path := NewTempDir(t)
-	repo, err := CloneRepo(origin.URL, path, true)
-	checkFatal(t, err)
-
-	wt, err := repo.repo.Worktree()
-	checkFatal(t, err)
-
-	files, err := wt.Filesystem.ReadDir("")
-	checkFatal(t, err)
-
-	if len(files) == 0 {
-		t.Errorf("Cloned repo should contain files")
-	}
-}
-
-func TestRepoEmpty(t *testing.T) {
-	repo := NewRepoEmpty(t)
-
-	wt, err := repo.Repo.Worktree()
-	checkFatal(t, err)
-
-	status, err := wt.Status()
-	checkFatal(t, err)
-	if !status.IsClean() {
-		t.Errorf("Empty repo should be clean")
-	}
-}
-
-func TestRepoWithUntrackedFile(t *testing.T) {
-	repo := NewRepoWithUntracked(t)
-
-	wt, err := repo.Repo.Worktree()
-	checkFatal(t, err)
-
-	status, err := wt.Status()
-	checkFatal(t, err)
-	if status.IsClean() {
-		t.Errorf("Repo with untracked file should not be clean")
-	}
-
-	// TODO: remove magic strings
-	if !status.IsUntracked("README") {
-		t.Errorf("New file should be untracked")
-	}
-}
-
-func TestRepoWithStagedFile(t *testing.T) {
-	repo := NewRepoWithStaged(t)
-
-	wt, err := repo.Repo.Worktree()
-	checkFatal(t, err)
-
-	status, err := wt.Status()
-	checkFatal(t, err)
-	if status.IsClean() {
-		t.Errorf("Repo with staged file should not be clean")
-	}
-
-	if status.IsUntracked("README") {
-		t.Errorf("Staged file should not be untracked")
-	}
-}
-
-func TestRepoWithSingleCommit(t *testing.T) {
-	repo := NewRepoWithCommit(t)
-
-	wt, err := repo.Repo.Worktree()
-	checkFatal(t, err)
-
-	status, err := wt.Status()
-	checkFatal(t, err)
-	if !status.IsClean() {
-		t.Errorf("Repo with committed file should be clean")
-	}
-
-	if status.IsUntracked("README") {
-		t.Errorf("Committed file should not be untracked")
-	}
-}
-func TestStatusWithModifiedFile(t *testing.T) {
-	//todo modified but not staged
-}
-
-func TestStatusWithUntrackedButIgnoredFile(t *testing.T) {
-	//todo
-}

new/status.go

@@ -1,6 +1,9 @@
 package new
 
 import (
+	"sort"
+	"strings"
+
 	"github.com/go-git/go-git/v5"
 	"github.com/go-git/go-git/v5/plumbing"
 	"github.com/pkg/errors"
@@ -9,21 +12,14 @@ import (
 type RepoStatus struct {
 	HasUntrackedFiles     bool
 	HasUncommittedChanges bool
-	Branches              map[string]*BranchStatus
+	Branches              []*BranchStatus
 }
 
 type BranchStatus struct {
 	Name      string
-	Upstream  *Upstream
+	Upstream  string
 	NeedsPull bool
 	NeedsPush bool
-	Ahead     int
-	Behind    int
-}
-
-type Upstream struct {
-	Remote string
-	Branch string
 }
 
 func (r *Repo) LoadStatus() error {
@@ -37,10 +33,10 @@ func (r *Repo) LoadStatus() error {
 		return errors.Wrap(err, "Failed getting worktree status")
 	}
 
-	r.Status.HasUncommittedChanges = !status.IsClean()
+	r.Status.HasUncommittedChanges = hasUncommitted(status)
 	r.Status.HasUntrackedFiles = hasUntracked(status)
 
-	err = r.LoadBranchesStatus()
+	err = r.loadBranchesStatus()
 	if err != nil {
 		return err
 	}
@@ -51,15 +47,32 @@ func (r *Repo) LoadStatus() error {
 // hasUntracked returns true if there are any untracked files in the worktree
 func hasUntracked(status git.Status) bool {
 	for _, fs := range status {
-		if fs.Worktree == git.Untracked {
+		if fs.Worktree == git.Untracked || fs.Staging == git.Untracked {
 			return true
 		}
 	}
-	return false
 
+	return false
 }
 
-func (r *Repo) LoadBranchesStatus() error {
+// hasUncommitted returns true if there are any uncommitted (but tracked) files in the worktree
+func hasUncommitted(status git.Status) bool {
+	// If repo is clean it means every file in worktree and staging has 'Unmodified' state
+	if status.IsClean() {
+		return false
+	}
+
+	// If repo is not clean, check if any file has state different than 'Untracked' - it means they are tracked and have uncommitted modifications
+	for _, fs := range status {
+		if fs.Worktree != git.Untracked || fs.Staging != git.Untracked {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (r *Repo) loadBranchesStatus() error {
 	iter, err := r.repo.Branches()
 	if err != nil {
 		return errors.Wrap(err, "Failed getting branches iterator")
@@ -71,60 +84,132 @@ func (r *Repo) LoadBranchesStatus() error {
 			return err
 		}
 
-		r.Status.Branches[bs.Name] = bs
+		r.Status.Branches = append(r.Status.Branches, bs)
 		return nil
 	})
 	if err != nil {
 		return errors.Wrap(err, "Failed iterating over branches")
 	}
 
+	// Sort branches by name. It's useful to have them sorted for printing and testing.
+	sort.Slice(r.Status.Branches, func(i, j int) bool {
+		return strings.Compare(r.Status.Branches[i].Name, r.Status.Branches[j].Name) < 0
+	})
 	return nil
 }
 
 func (r *Repo) newBranchStatus(branch string) (*BranchStatus, error) {
+	bs := &BranchStatus{
+		Name: branch,
+	}
+
 	upstream, err := r.upstream(branch)
 	if err != nil {
 		return nil, err
 	}
 
-	return &BranchStatus{
-		Name:     branch,
-		Upstream: upstream,
-	}, nil
+	if upstream == "" {
+		return bs, nil
+	}
+
+	needsPull, needsPush, err := r.needsPullOrPush(branch, upstream)
+	if err != nil {
+		return nil, err
+	}
+
+	bs.Upstream = upstream
+	bs.NeedsPush = needsPush
+	bs.NeedsPull = needsPull
+
+	return bs, nil
 }
 
 // upstream finds if a given branch tracks an upstream.
-// Returns found upstream or nil if branch doesn't track an upstream.
+// Returns found upstream branch name (eg, origin/master) or empty string if branch doesn't track an upstream.
 //
 // Information about upstream is taken from .git/config file.
 // If a branch has an upstream, there's a [branch] section in the file with two fields:
 // "remote" - name of the remote containing upstreamn branch (or "." if upstream is a local branch)
 // "merge" - full ref name of the upstream branch (eg, ref/heads/master)
-func (r *Repo) upstream(branch string) (*Upstream, error) {
+func (r *Repo) upstream(branch string) (string, error) {
 	cfg, err := r.repo.Config()
 	if err != nil {
-		return nil, errors.Wrap(err, "Failed getting repo config")
+		return "", errors.Wrap(err, "Failed getting repo config")
 	}
 
-	// Find our branch in "branch" config sections
+	// Check if our branch exists in "branch" config sections. If not, it doesn't have an upstream configured.
 	bcfg := cfg.Branches[branch]
 	if bcfg == nil {
-		return nil, nil
+		return "", nil
 	}
 
 	remote := bcfg.Remote
 	if remote == "" {
-		return nil, nil
+		return "", nil
 	}
 
-	// TODO: check if this should be short or full ref name
 	merge := bcfg.Merge.Short()
 	if merge == "" {
-		return nil, nil
+		return "", nil
+	}
+	return remote + "/" + merge, nil
+}
+
+func (r *Repo) needsPullOrPush(localBranch string, upstreamBranch string) (needsPull bool, needsPush bool, err error) {
+	localHash, err := r.repo.ResolveRevision(plumbing.Revision(localBranch))
+	if err != nil {
+		return false, false, errors.Wrapf(err, "Failed resolving revision %s", localBranch)
 	}
 
-	return &Upstream{
-		Remote: remote,
-		Branch: merge,
-	}, nil
+	upstreamHash, err := r.repo.ResolveRevision(plumbing.Revision(upstreamBranch))
+	if err != nil {
+		return false, false, errors.Wrapf(err, "Failed resolving revision %s", upstreamBranch)
+	}
+
+	localCommit, err := r.repo.CommitObject(*localHash)
+	if err != nil {
+		return false, false, errors.Wrapf(err, "Failed finding a commit for hash %s", localHash.String())
+	}
+
+	upstreamCommit, err := r.repo.CommitObject(*upstreamHash)
+	if err != nil {
+		return false, false, errors.Wrapf(err, "Failed finding a commit for hash %s", upstreamHash.String())
+	}
+
+	// If local branch hash is the same as upstream, it means there is no difference between local and upstream
+	if *localHash == *upstreamHash {
+		return false, false, nil
+	}
+
+	commons, err := localCommit.MergeBase(upstreamCommit)
+	if err != nil {
+		return false, false, errors.Wrapf(err, "Failed finding common ancestors for branches %s & %s", localBranch, upstreamBranch)
+	}
+
+	if len(commons) == 0 {
+		// TODO: No common ancestors. This should be an error
+		return false, false, nil
+	}
+
+	if len(commons) > 1 {
+		// TODO: multiple best ancestors. How to handle this?
+		return false, false, nil
+	}
+
+	mergeBase := commons[0]
+
+	// If merge base is the same as upstream branch, local branch is ahead and push is needed
+	// If merge base is the same as local branch, local branch is behind and pull is needed
+	// If merge base is something else, branches have diverged and merge is needed (both pull and push)
+	// ref: https://stackoverflow.com/a/17723781/1085632
+
+	if mergeBase.Hash == *upstreamHash {
+		return false, true, nil
+	}
+
+	if mergeBase.Hash == *localHash {
+		return true, false, nil
+	}
+
+	return true, true, nil
 }

new/status_test.go

@@ -1,44 +1,149 @@
 package new
 
-import "testing"
+import (
+	"reflect"
+	"testing"
+)
 
-func TestBranchStatusLocal(t *testing.T) {
-	tr := NewRepoWithCommit(t)
-	tr.NewBranch("branch")
-
-	repo, err := OpenRepo(tr.Path)
-	checkFatal(t, err)
-
-	err = repo.LoadStatus()
-	checkFatal(t, err)
-
-	if repo.Status.Branches["master"].Upstream != nil {
-		t.Errorf("'master' branch should not have an upstream")
+func TestStatus(t *testing.T) {
+	var tests = []struct {
+		makeTestRepo func(*testing.T) *TestRepo
+		want         *RepoStatus
+	}{
+		{NewRepoEmpty, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches:              nil,
+		}},
+		{NewRepoWithUntracked, &RepoStatus{
+			HasUntrackedFiles:     true,
+			HasUncommittedChanges: false,
+			Branches:              nil,
+		}},
+		{NewRepoWithStaged, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: true,
+			Branches:              nil,
+		}},
+		{NewRepoWithCommit, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches: []*BranchStatus{
+				{
+					Name:      "master",
+					Upstream:  "",
+					NeedsPull: false,
+					NeedsPush: false,
+				},
+			},
+		}},
+		{NewRepoWithModified, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: true,
+			Branches: []*BranchStatus{
+				{
+					Name:      "master",
+					Upstream:  "",
+					NeedsPull: false,
+					NeedsPush: false,
+				},
+			},
+		}},
+		{NewRepoWithIgnored, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches: []*BranchStatus{
+				{
+					Name:      "master",
+					Upstream:  "",
+					NeedsPull: false,
+					NeedsPush: false,
+				},
+			},
+		}},
+		{NewRepoWithLocalBranch, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches: []*BranchStatus{
+				{
+					Name:      "local",
+					Upstream:  "",
+					NeedsPull: false,
+					NeedsPush: false,
+				}, {
+					Name:      "master",
+					Upstream:  "",
+					NeedsPull: false,
+					NeedsPush: false,
+				},
+			},
+		}},
+		{NewRepoWithClonedBranch, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches: []*BranchStatus{
+				{
+					Name:      "local",
+					Upstream:  "",
+					NeedsPull: false,
+					NeedsPush: false,
+				}, {
+					Name:      "master",
+					Upstream:  "origin/master",
+					NeedsPull: false,
+					NeedsPush: false,
+				},
+			},
+		}},
+		{NewRepoWithBranchAhead, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches: []*BranchStatus{
+				{
+					Name:      "master",
+					Upstream:  "origin/master",
+					NeedsPull: false,
+					NeedsPush: true,
+				},
+			},
+		}},
+		{NewRepoWithBranchBehind, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches: []*BranchStatus{
+				{
+					Name:      "master",
+					Upstream:  "origin/master",
+					NeedsPull: true,
+					NeedsPush: false,
+				},
+			},
+		}},
+		{NewRepoWithBranchAheadAndBehind, &RepoStatus{
+			HasUntrackedFiles:     false,
+			HasUncommittedChanges: false,
+			Branches: []*BranchStatus{
+				{
+					Name:      "master",
+					Upstream:  "origin/master",
+					NeedsPull: true,
+					NeedsPush: true,
+				},
+			},
+		}},
 	}
 
-	if repo.Status.Branches["branch"].Upstream != nil {
-		t.Errorf("'branch' branch should not have an upstream")
+	for _, test := range tests {
+		tr := test.makeTestRepo(t)
+
+		repo, err := OpenRepo(tr.Path)
+		checkFatal(t, err)
+
+		err = repo.LoadStatus()
+		checkFatal(t, err)
+
+		if !reflect.DeepEqual(repo.Status, test.want) {
+			t.Errorf("Wrong repo status, got: %+v; want: %+v", repo.Status, test.want)
+		}
 	}
 }
-
-func TestBranchStatusCloned(t *testing.T) {
-	origin := NewRepoWithCommit(t)
-
-	clone := origin.Clone()
-	clone.NewBranch("local")
-
-	repo, err := OpenRepo(clone.Path)
-	checkFatal(t, err)
-
-	err = repo.LoadStatus()
-	checkFatal(t, err)
-
-	if repo.Status.Branches["master"].Upstream == nil {
-		t.Errorf("'master' branch should have an upstream")
-	}
-
-	if repo.Status.Branches["local"].Upstream != nil {
-		t.Errorf("'local' branch should not have an upstream")
-	}
-
-}