Refactor Repo code and tests

- Add Repo as a separate type
- Merge together repo status and branches status
- Load status on repo opening

pkg/status.go

@@ -8,20 +8,32 @@ import (
 type RepoStatus struct {
 	HasUntrackedFiles     bool
 	HasUncommittedChanges bool
-	BranchStatuses        map[string]BranchStatus
+	Branches              map[string]BranchStatus
 }
 
-func NewRepoStatus(path string) (RepoStatus, error) {
-	var status RepoStatus
+type BranchStatus struct {
+	Name        string
+	IsRemote    bool
+	HasUpstream bool
+	NeedsPull   bool
+	NeedsPush   bool
+	Ahead       int
+	Behind      int
+}
 
-	repo, err := git.OpenRepository(path)
+func loadStatus(r *git.Repository) (*RepoStatus, error) {
+	entries, err := statusEntries(r)
 	if err != nil {
-		return status, errors.Wrap(err, "Failed opening repository")
+		return nil, err
 	}
 
-	entries, err := statusEntries(repo)
+	branches, err := branches(r)
 	if err != nil {
-		return status, errors.Wrap(err, "Failed getting repository status")
+		return nil, err
+	}
+
+	status := &RepoStatus{
+		Branches: branches,
 	}
 
 	for _, entry := range entries {
@@ -33,29 +45,23 @@ func NewRepoStatus(path string) (RepoStatus, error) {
 		}
 	}
 
-	branchStatuses, err := Branches(repo)
-	if err != nil {
-		return status, errors.Wrap(err, "Failed getting branches statuses")
-	}
-	status.BranchStatuses = branchStatuses
-
 	return status, nil
 }
 
-func statusEntries(repo *git.Repository) ([]git.StatusEntry, error) {
+func statusEntries(r *git.Repository) ([]git.StatusEntry, error) {
 	opts := &git.StatusOptions{
 		Show:  git.StatusShowIndexAndWorkdir,
 		Flags: git.StatusOptIncludeUntracked,
 	}
 
-	status, err := repo.StatusList(opts)
+	status, err := r.StatusList(opts)
 	if err != nil {
-		return nil, errors.Wrap(err, "Failed getting repository status")
+		return nil, errors.Wrap(err, "Failed getting repository status list")
 	}
 
 	entryCount, err := status.EntryCount()
 	if err != nil {
-		return nil, errors.Wrap(err, "Failed getting repository status count")
+		return nil, errors.Wrap(err, "Failed getting repository status list count")
 	}
 
 	var entries []git.StatusEntry
@@ -70,3 +76,78 @@ func statusEntries(repo *git.Repository) ([]git.StatusEntry, error) {
 
 	return entries, nil
 }
+
+func branches(r *git.Repository) (map[string]BranchStatus, error) {
+	iter, err := r.NewBranchIterator(git.BranchAll)
+	if err != nil {
+		return nil, errors.Wrap(err, "Failed creating branch iterator")
+	}
+
+	var branches []*git.Branch
+	err = iter.ForEach(func(branch *git.Branch, branchType git.BranchType) error {
+		branches = append(branches, branch)
+		return nil
+	})
+
+	if err != nil {
+		return nil, errors.Wrap(err, "Failed iterating over branches")
+	}
+
+	statuses := make(map[string]BranchStatus)
+	for _, branch := range branches {
+		status, err := branchStatus(branch)
+		if err != nil {
+			// TODO: Handle error. We should tell user that we couldn't read status of that branch but probably shouldn't exit
+			continue
+		}
+		statuses[status.Name] = status
+	}
+
+	return statuses, nil
+}
+
+func branchStatus(branch *git.Branch) (BranchStatus, error) {
+	var status BranchStatus
+
+	name, err := branch.Name()
+	if err != nil {
+		return status, errors.Wrap(err, "Failed getting branch name")
+	}
+	status.Name = name
+
+	// If branch is a remote one, return immediately. Upstream can only be found for local branches.
+	if branch.IsRemote() {
+		status.IsRemote = true
+		return status, nil
+	}
+
+	upstream, err := branch.Upstream()
+	if err != nil && !git.IsErrorCode(err, git.ErrNotFound) {
+		return status, errors.Wrap(err, "Failed getting branch upstream")
+	}
+
+	// If there's no upstream, return immediately. Ahead/Behind can only be found when upstream exists.
+	if upstream == nil {
+		return status, nil
+	}
+
+	status.HasUpstream = true
+
+	ahead, behind, err := branch.Owner().AheadBehind(branch.Target(), upstream.Target())
+	if err != nil {
+		return status, errors.Wrap(err, "Failed getting ahead/behind information")
+	}
+
+	status.Ahead = ahead
+	status.Behind = behind
+
+	if ahead > 0 {
+		status.NeedsPush = true
+	}
+
+	if behind > 0 {
+		status.NeedsPull = true
+	}
+
+	return status, nil
+}