pipeline: add Pipeline.Branch to define where to find CommitsHead

Signed-off-by: Robert Lin <robert@bobheadxi.dev>

cmd/hercules/root.go

@@ -24,7 +24,7 @@ import (
 	"github.com/spf13/pflag"
 	"golang.org/x/crypto/ssh/terminal"
 	progress "gopkg.in/cheggaaa/pb.v1"
-	"gopkg.in/src-d/go-billy-siva.v4"
+	sivafs "gopkg.in/src-d/go-billy-siva.v4"
 	"gopkg.in/src-d/go-billy.v4/memfs"
 	"gopkg.in/src-d/go-billy.v4/osfs"
 	"gopkg.in/src-d/go-git.v4"
@@ -190,6 +190,7 @@ targets can be added using the --plugin system.`,
 		}
 		firstParent := getBool("first-parent")
 		commitsFile := getString("commits")
+		branch := getString("branch")
 		head := getBool("head")
 		protobuf := getBool("pb")
 		profile := getBool("profile")
@@ -219,6 +220,7 @@ targets can be added using the --plugin system.`,
 
 		// core logic
 		pipeline := hercules.NewPipeline(repository)
+		pipeline.Branch = branch
 		pipeline.SetFeaturesFromFlags()
 		var bar *progress.ProgressBar
 		if !disableStatus {
@@ -484,28 +486,33 @@ var cmdlineDeployed map[string]*bool
 func init() {
 	loadPlugins()
 	rootFlags := rootCmd.Flags()
+
+	// commits flags
 	rootFlags.String("commits", "", "Path to the text file with the "+
 		"commit history to follow instead of the default 'git log'. "+
 		"The format is the list of hashes, each hash on a "+
 		"separate line. The first hash is the root.")
-	err := rootCmd.MarkFlagFilename("commits")
-	if err != nil {
+	if err := rootCmd.MarkFlagFilename("commits"); err != nil {
 		panic(err)
 	}
 	hercules.PathifyFlagValue(rootFlags.Lookup("commits"))
+	rootFlags.String("branch", "", "Specify a branch to analyze.")
 	rootFlags.Bool("head", false, "Analyze only the latest commit.")
 	rootFlags.Bool("first-parent", false, "Follow only the first parent in the commit history - "+
 		"\"git log --first-parent\".")
+
+	// output flags
 	rootFlags.Bool("pb", false, "The output format will be Protocol Buffers instead of YAML.")
 	rootFlags.Bool("quiet", !terminal.IsTerminal(int(os.Stdin.Fd())),
 		"Do not print status updates to stderr.")
 	rootFlags.Bool("profile", false, "Collect the profile to hercules.pprof.")
 	rootFlags.String("ssh-identity", "", "Path to SSH identity file (e.g., ~/.ssh/id_rsa) to clone from an SSH remote.")
-	err = rootCmd.MarkFlagFilename("ssh-identity")
-	if err != nil {
+	if err := rootCmd.MarkFlagFilename("ssh-identity"); err != nil {
 		panic(err)
 	}
 	hercules.PathifyFlagValue(rootFlags.Lookup("ssh-identity"))
+
+	// register all flags
 	cmdlineFacts, cmdlineDeployed = hercules.Registry.AddFlags(rootFlags)
 	rootCmd.SetUsageFunc(formatUsage)
 	rootCmd.AddCommand(versionCmd)

internal/core/pipeline.go

@@ -272,6 +272,9 @@ type Pipeline struct {
 	// PrintActions indicates whether to print the taken actions during the execution.
 	PrintActions bool
 
+	// Branch used for pipeline.HeadCommit. Leave blank to use HEAD.
+	Branch string
+
 	// Repository points to the analysed Git repository struct from go-git.
 	repository *git.Repository
 
@@ -484,34 +487,59 @@ func (pipeline *Pipeline) Commits(firstParent bool) ([]*object.Commit, error) {
 // HeadCommit returns the latest commit in the repository (HEAD).
 func (pipeline *Pipeline) HeadCommit() ([]*object.Commit, error) {
 	repository := pipeline.repository
-	head, err := repository.Head()
-	if err == plumbing.ErrReferenceNotFound {
-		refs, errr := repository.References()
-		if errr != nil {
-			return nil, errors.Wrap(errr, "unable to list the references")
-		}
-		var refnames []string
-		refByName := map[string]*plumbing.Reference{}
-		err = refs.ForEach(func(ref *plumbing.Reference) error {
-			refname := ref.Name().String()
-			refnames = append(refnames, refname)
-			refByName[refname] = ref
-			if strings.HasPrefix(refname, "refs/heads/HEAD/") {
+
+	var head *plumbing.Reference
+	if pipeline.Branch != "" {
+		pipeline.l.Infof("querying for head of branch %s", pipeline.Branch)
+		branch := plumbing.NewBranchReferenceName(pipeline.Branch)
+		iter, err := repository.Branches()
+		if err != nil {
+			return nil, errors.Wrap(err, "unable to list branches")
+		}
+		if err := iter.ForEach(func(ref *plumbing.Reference) error {
+			pipeline.l.Info(ref.Name())
+			if ref.Name() == branch {
 				head = ref
 				return storer.ErrStop
 			}
 			return nil
-		})
-		if head == nil {
-			sort.Strings(refnames)
-			headName := refnames[len(refnames)-1]
-			pipeline.l.Warnf("could not determine the HEAD, falling back to %s", headName)
-			head = refByName[headName]
+		}); err != nil {
+			return nil, errors.Wrap(err, "unable to find branch head")
+		}
+	} else {
+		var err error
+		head, err = repository.Head()
+		if err == plumbing.ErrReferenceNotFound {
+			refs, errr := repository.References()
+			if errr != nil {
+				return nil, errors.Wrap(errr, "unable to list the references")
+			}
+			var refnames []string
+			refByName := map[string]*plumbing.Reference{}
+			err = refs.ForEach(func(ref *plumbing.Reference) error {
+				refname := ref.Name().String()
+				refnames = append(refnames, refname)
+				refByName[refname] = ref
+				if strings.HasPrefix(refname, "refs/heads/HEAD/") {
+					head = ref
+					return storer.ErrStop
+				}
+				return nil
+			})
+			if head == nil {
+				sort.Strings(refnames)
+				headName := refnames[len(refnames)-1]
+				pipeline.l.Warnf("could not determine the HEAD, falling back to %s", headName)
+				head = refByName[headName]
+			}
+		} else if err != nil {
+			return nil, errors.Wrap(err, "unable to find the head reference")
 		}
 	}
 	if head == nil {
-		return nil, errors.Wrap(err, "unable to find the head reference")
+		return nil, errors.New("unable to find the head reference")
 	}
+
 	commit, err := repository.CommitObject(head.Hash())
 	if err != nil {
 		return nil, err

internal/core/pipeline_test.go

@@ -442,13 +442,30 @@ func TestPipelineCommitsFirstParent(t *testing.T) {
 }
 
 func TestPipelineHeadCommit(t *testing.T) {
-	pipeline := NewPipeline(test.Repository)
-	commits, err := pipeline.HeadCommit()
-	assert.NoError(t, err)
-	assert.Len(t, commits, 1)
-	assert.True(t, len(commits[0].ParentHashes) > 0)
-	head, _ := test.Repository.Head()
-	assert.Equal(t, head.Hash(), commits[0].Hash)
+	t.Run("default", func(t *testing.T) {
+		pipeline := NewPipeline(test.Repository)
+		commits, err := pipeline.HeadCommit()
+		assert.NoError(t, err)
+		assert.Len(t, commits, 1)
+		assert.True(t, len(commits[0].ParentHashes) > 0)
+		head, _ := test.Repository.Head()
+		assert.Equal(t, head.Hash(), commits[0].Hash)
+	})
+	t.Run("with branch specified", func(t *testing.T) {
+		testBranch := "test-branch"
+		repo, out := test.NewInMemRepository(&test.InMemRepositoryOptions{
+			CreateBranch: testBranch,
+		})
+		assert.False(t, out.CreatedBranchHash.IsZero())
+
+		pipeline := NewPipeline(repo)
+		pipeline.Branch = testBranch
+		commits, err := pipeline.HeadCommit()
+		assert.NoError(t, err)
+		assert.Len(t, commits, 1)
+		assert.True(t, len(commits[0].ParentHashes) > 0)
+		assert.Equal(t, out.CreatedBranchHash, commits[0].Hash)
+	})
 }
 
 func TestLoadCommitsFromFile(t *testing.T) {

internal/test/repository.go

@@ -1,11 +1,14 @@
 package test
 
 import (
+	"fmt"
 	"io"
 	"io/ioutil"
 	"os"
 	"path"
+	"time"
 
+	"gopkg.in/src-d/go-billy.v4/memfs"
 	"gopkg.in/src-d/go-git.v4"
 	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
@@ -75,3 +78,59 @@ func init() {
 		panic(err)
 	}
 }
+
+// InMemRepositoryOptions declares config for NewInMemRepository
+type InMemRepositoryOptions struct {
+	CreateBranch string
+}
+
+// InMemRepositoryOutput provides output from options provided in InMemRepositoryOptions
+type InMemRepositoryOutput struct {
+	CreatedBranchHash plumbing.Hash
+}
+
+// NewInMemRepository initializes a new in-memory repository
+func NewInMemRepository(opts *InMemRepositoryOptions) (*git.Repository, InMemRepositoryOutput) {
+	var out InMemRepositoryOutput
+
+	repo, err := git.Clone(memory.NewStorage(), memfs.New(), &git.CloneOptions{
+		URL: "https://github.com/src-d/hercules",
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	if opts != nil && opts.CreateBranch != "" {
+		t, err := repo.Worktree()
+		if err != nil {
+			panic(err)
+		}
+		if err := t.Checkout(&git.CheckoutOptions{
+			Branch: plumbing.NewBranchReferenceName(opts.CreateBranch),
+			Create: true,
+		}); err != nil {
+			panic(err)
+		}
+		if _, err := t.Add("."); err != nil {
+			panic(err)
+		}
+		out.CreatedBranchHash, err = t.Commit(
+			fmt.Sprintf("test commit on %s", opts.CreateBranch),
+			&git.CommitOptions{
+				Author: &object.Signature{Name: "bobheadxi", Email: "bobheadxi@email.com", When: time.Now()},
+			},
+		)
+		if err != nil {
+			panic(err)
+		}
+
+		// check out master again
+		if err := t.Checkout(&git.CheckoutOptions{
+			Branch: plumbing.NewBranchReferenceName("master"),
+		}); err != nil {
+			panic(err)
+		}
+	}
+
+	return repo, out
+}