Repository: palantir/policy-bot (go)

Issue:
support predicate for source branch
It seems like it would be useful to have an analogous version of `targets_branch` for specifying the source branch. For my particular use case, I would like to have a policy where pull requests to the `master` branch can only come from the `develop` branch.

Pull request description:
Add from_branch predicate
Fixes #241.

This adds a new predicate called `source_branch` that's analogous to `targets_branch`, but for the PR's source branch. I'm very open to different names for this option, but I haven't come up with something I love.

I changed the existing `targets_branch` tests to also run the patterns against the source branch. I also added tests for making sure both branches are correctly set when creating the `GitHubContext` objects.

Patch:
diff --git a/README.md b/README.md
index ed92545..dd5e75b 100644
--- a/README.md
+++ b/README.md
@@ -171,6 +171,12 @@ if:
   targets_branch:
     pattern: "^(master|regexPattern)$"
 
+  # "from_branch" is satisfied if the source branch of the pull request
+  # matches the regular expression. Note that source branches from forks will
+  # have the pattern "repo_owner:branch_name"
+  from_branch:
+    pattern: "^(master|regexPattern)$"
+
   # "modified_lines" is satisfied if the number of lines added or deleted by
   # the pull request matches any of the listed conditions. Each expression is
   # an operator (one of '<' or '>'), an optional space, and a number.
diff --git a/policy/approval/predicate.go b/policy/approval/predicate.go
index a8476c1..e31d9a2 100644
--- a/policy/approval/predicate.go
+++ b/policy/approval/predicate.go
@@ -28,6 +28,7 @@ type Predicates struct {
 	AuthorIsOnlyContributor *predicate.AuthorIsOnlyContributor `yaml:"author_is_only_contributor"`
 
 	TargetsBranch *predicate.TargetsBranch `yaml:"targets_branch"`
+	FromBranch    *predicate.FromBranch    `yaml:"from_branch"`
 
 	ModifiedLines *predicate.ModifiedLines `yaml:"modified_lines"`
 
@@ -62,6 +63,10 @@ func (p *Predicates) Predicates() []predicate.Predicate {
 	if p.TargetsBranch != nil {
 		ps = append(ps, predicate.Predicate(p.TargetsBranch))
 	}
+	if p.FromBranch != nil {
+		ps = append(ps, predicate.Predicate(p.FromBranch))
+	}
+
 	if p.ModifiedLines != nil {
 		ps = append(ps, predicate.Predicate(p.ModifiedLines))
 	}
diff --git a/policy/predicate/branch.go b/policy/predicate/branch.go
index 0491fc0..1ed5953 100644
--- a/policy/predicate/branch.go
+++ b/policy/predicate/branch.go
@@ -39,3 +39,21 @@ func (pred *TargetsBranch) Evaluate(ctx context.Context, prctx pull.Context) (bo
 
 	return matches, desc, nil
 }
+
+type FromBranch struct {
+	Pattern common.Regexp `yaml:"pattern"`
+}
+
+var _ Predicate = &FromBranch{}
+
+func (pred *FromBranch) Evaluate(ctx context.Context, prctx pull.Context) (bool, string, error) {
+	_, sourceBranchName := prctx.Branches()
+	matches := pred.Pattern.Matches(sourceBranchName)
+
+	desc := ""
+	if !matches {
+		desc = fmt.Sprintf("Source branch %q does not match specified from_branch pattern %q", sourceBranchName, pred.Pattern)
+	}
+
+	return matches, desc, nil
+}


Test patch:
diff --git a/policy/predicate/branch_test.go b/policy/predicate/branch_test.go
index 35e621b..ae7a617 100644
--- a/policy/predicate/branch_test.go
+++ b/policy/predicate/branch_test.go
@@ -22,104 +22,98 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/palantir/policy-bot/policy/common"
-	"github.com/palantir/policy-bot/pull"
 	"github.com/palantir/policy-bot/pull/pulltest"
 )
 
-func TestTargetsBranch(t *testing.T) {
-	p := &TargetsBranch{
-		Pattern: common.NewCompiledRegexp(regexp.MustCompile("^master$")),
-	}
-
-	runTargetsTestCase(t, p, []targetsTestCase{
+func TestBranches(t *testing.T) {
+	runBranchesTestCase(t, "^master$", []branchesTestCase{
 		{
 			"simple match - master",
 			true,
-			&pulltest.Context{
-				BranchBaseName: "master",
-			},
+			"master",
 		},
 		{
 			"simple non match",
 			false,
-			&pulltest.Context{
-				BranchBaseName: "another-branch",
-			},
+			"another-branch",
 		},
 		{
 			"tests anchoring",
 			false,
-			&pulltest.Context{
-				BranchBaseName: "not-master",
-			},
+			"not-master",
 		},
 	})
 
-	pMatchAll := &TargetsBranch{
-		Pattern: common.NewCompiledRegexp(regexp.MustCompile(".*")),
-	}
-
-	runTargetsTestCase(t, pMatchAll, []targetsTestCase{
+	runBranchesTestCase(t, ".*", []branchesTestCase{
 		{
 			"matches all example 1",
 			true,
-			&pulltest.Context{
-				BranchBaseName: "master",
-			},
+			"master",
 		},
 		{
 			"matches all example 2",
 			true,
-			&pulltest.Context{
-				BranchBaseName: "another-one",
-			},
+			"another-one",
 		},
 	})
 
-	pRegex := &TargetsBranch{
-		Pattern: common.NewCompiledRegexp(regexp.MustCompile("(prod|staging)")),
-	}
-
-	runTargetsTestCase(t, pRegex, []targetsTestCase{
+	runBranchesTestCase(t, "(prod|staging)", []branchesTestCase{
 		{
 			"matches pattern - prod",
 			true,
-			&pulltest.Context{
-				BranchBaseName: "prod",
-			},
+			"prod",
 		},
 		{
 			"matches pattern - staging",
 			true,
-			&pulltest.Context{
-				BranchBaseName: "staging",
-			},
+			"staging",
 		},
 		{
 			"matches pattern - not-a-match",
 			false,
-			&pulltest.Context{
-				BranchBaseName: "not-a-match",
-			},
+			"not-a-match",
 		},
 	})
 }
 
 // TODO: generalize this and use it all our test cases
-type targetsTestCase struct {
-	name     string
-	expected bool
-	context  pull.Context
+type branchesTestCase struct {
+	name       string
+	expected   bool
+	branchName string
 }
 
-func runTargetsTestCase(t *testing.T, p Predicate, cases []targetsTestCase) {
+func runBranchesTestCase(t *testing.T, regex string, cases []branchesTestCase) {
 	ctx := context.Background()
 
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
-			ok, _, err := p.Evaluate(ctx, tc.context)
-			if assert.NoError(t, err, "evaluation failed") {
-				assert.Equal(t, tc.expected, ok, "predicate was not correct")
+
+		compiledRegexp := common.NewCompiledRegexp(regexp.MustCompile(regex))
+		targetsPredicate := &TargetsBranch{
+			Pattern: compiledRegexp,
+		}
+		fromPredicate := &FromBranch{
+			Pattern: compiledRegexp,
+		}
+
+		targetsContext := &pulltest.Context{
+			BranchBaseName: tc.branchName,
+		}
+		fromContext := &pulltest.Context{
+			BranchHeadName: tc.branchName,
+		}
+
+		t.Run(tc.name+" targets_branch", func(t *testing.T) {
+			ok, _, err := targetsPredicate.Evaluate(ctx, targetsContext)
+			if assert.NoError(t, err, "targets_branch predicate evaluation failed") {
+				assert.Equal(t, tc.expected, ok, "targets_branch predicate was not correct")
+			}
+		})
+
+		t.Run(tc.name+" from_branch", func(t *testing.T) {
+			ok, _, err := fromPredicate.Evaluate(ctx, fromContext)
+			if assert.NoError(t, err, "from_branch predicate evaluation failed") {
+				assert.Equal(t, tc.expected, ok, "from_branch predicate was not correct")
 			}
 		})
 	}
diff --git a/pull/github_test.go b/pull/github_test.go
index aa7c4cb..b4bed92 100644
--- a/pull/github_test.go
+++ b/pull/github_test.go
@@ -396,6 +396,35 @@ func TestIsOrgMember(t *testing.T) {
 	assert.Equal(t, 1, yesRule.Count, "cached membership was not used")
 }
 
+func TestBranches(t *testing.T) {
+	rp := &ResponsePlayer{}
+	ctx := makeContext(t, rp, nil)
+
+	base, head := ctx.Branches()
+	assert.Equal(t, "develop", base, "base branch was not correctly set")
+	assert.Equal(t, "test-branch", head, "head branch was not correctly set")
+}
+
+func TestCrossRepoBranches(t *testing.T) {
+	rp := &ResponsePlayer{}
+
+	// change the source repo to a forked repo
+	crossRepoPr := defaultTestPR()
+	crossRepoPr.Head.Repo = &github.Repository{
+		ID: github.Int64(12345),
+		Owner: &github.User{
+			Login: github.String("testorg2"),
+		},
+		Name: github.String("testrepofork"),
+	}
+
+	ctx := makeContext(t, rp, crossRepoPr)
+
+	base, head := ctx.Branches()
+	assert.Equal(t, "develop", base, "cross-repo base branch was not correctly set")
+	assert.Equal(t, "testorg2:test-branch", head, "cross-repo head branch was not correctly set")
+}
+
 func makeContext(t *testing.T, rp *ResponsePlayer, pr *github.PullRequest) Context {
 	ctx := context.Background()
 	client := github.NewClient(&http.Client{Transport: rp})