From 744ba8280e18a7088791738f24cb51dc537cfcd5 Mon Sep 17 00:00:00 2001
From: Billy Keyes <bluekeyes@gmail.com>
Date: Mon, 25 Mar 2019 17:21:40 -0700
Subject: [PATCH] Add simple test for PushRestrictionMerger

---
 bulldozer/merge_test.go | 46 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/bulldozer/merge_test.go b/bulldozer/merge_test.go
index 3f12b2b01..e34b5419e 100644
--- a/bulldozer/merge_test.go
+++ b/bulldozer/merge_test.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/google/go-github/github"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -25,6 +26,24 @@ import (
 	"github.com/palantir/bulldozer/pull/pulltest"
 )
 
+type MockMerger struct {
+	MergeCount int
+	MergeError error
+
+	DeleteCount int
+	DeleteError error
+}
+
+func (m *MockMerger) Merge(ctx context.Context, pullCtx pull.Context, message string, options *github.PullRequestOptions) (string, error) {
+	m.MergeCount++
+	return "deadbeef", m.MergeError
+}
+
+func (m *MockMerger) DeleteHead(ctx context.Context, pullCtx pull.Context) error {
+	m.DeleteCount++
+	return m.DeleteError
+}
+
 func TestCalculateCommitTitle(t *testing.T) {
 	defaultPullContext := &pulltest.MockPullContext{
 		NumberValue: 12,
@@ -78,3 +97,30 @@ func TestCalculateCommitTitle(t *testing.T) {
 		})
 	}
 }
+
+func TestPushRestrictionMerger(t *testing.T) {
+	normal := &MockMerger{}
+	restricted := &MockMerger{}
+	merger := NewPushRestrictionMerger(normal, restricted)
+
+	ctx := context.Background()
+	pullCtx := &pulltest.MockPullContext{}
+
+	_, _ = merger.Merge(ctx, pullCtx, "", nil)
+	assert.Equal(t, 1, normal.MergeCount, "normal merge was not called")
+	assert.Equal(t, 0, restricted.MergeCount, "restricted merge was incorrectly called")
+
+	_ = merger.DeleteHead(ctx, pullCtx)
+	assert.Equal(t, 1, normal.DeleteCount, "normal delete was not called")
+	assert.Equal(t, 0, restricted.DeleteCount, "restricted delete was incorrectly called")
+
+	pullCtx.PushRestrictionsValue = true
+
+	_, _ = merger.Merge(ctx, pullCtx, "", nil)
+	assert.Equal(t, 1, normal.MergeCount, "normal merge was incorrectly called")
+	assert.Equal(t, 1, restricted.MergeCount, "restricted merge was not called")
+
+	_ = merger.DeleteHead(ctx, pullCtx)
+	assert.Equal(t, 1, normal.DeleteCount, "normal delete was incorrectly called")
+	assert.Equal(t, 1, restricted.DeleteCount, "restricted delete was not called")
+}