From 8681e7a7b62381d6c097572ef6810fa1fecbf939 Mon Sep 17 00:00:00 2001
From: Corey Johnson <probablycorey@gmail.com>
Date: Tue, 5 May 2020 15:09:02 -0700
Subject: [PATCH] Make squash and rebase work

---
 api/queries_pr.go  | 13 +++++++++++++
 command/pr.go      | 31 +++++++++++++++++++++++++++----
 command/pr_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/api/queries_pr.go b/api/queries_pr.go
index 6da60ce31..ffd31c75b 100644
--- a/api/queries_pr.go
+++ b/api/queries_pr.go
@@ -800,6 +800,18 @@ func PullRequestReopen(client *Client, repo ghrepo.Interface, pr *PullRequest) e
 }
 
 func PullRequestMerge(client *Client, repo ghrepo.Interface, pr *PullRequest) error {
+	return merge(client, repo, pr, githubv4.PullRequestMergeMethodMerge)
+}
+
+func PullRequestRebase(client *Client, repo ghrepo.Interface, pr *PullRequest) error {
+	return merge(client, repo, pr, githubv4.PullRequestMergeMethodRebase)
+}
+
+func PullRequestSquash(client *Client, repo ghrepo.Interface, pr *PullRequest) error {
+	return merge(client, repo, pr, githubv4.PullRequestMergeMethodSquash)
+}
+
+func merge(client *Client, repo ghrepo.Interface, pr *PullRequest, mergeMethod githubv4.PullRequestMergeMethod) error {
 	var mutation struct {
 		ReopenPullRequest struct {
 			PullRequest struct {
@@ -810,6 +822,7 @@ func PullRequestMerge(client *Client, repo ghrepo.Interface, pr *PullRequest) er
 
 	input := githubv4.MergePullRequestInput{
 		PullRequestID: pr.ID,
+		MergeMethod:   &mergeMethod,
 	}
 
 	v4 := githubv4.NewClient(client.http)
diff --git a/command/pr.go b/command/pr.go
index 9e73e4709..4ff1a2ad3 100644
--- a/command/pr.go
+++ b/command/pr.go
@@ -25,6 +25,9 @@ func init() {
 	prCmd.AddCommand(prCloseCmd)
 	prCmd.AddCommand(prReopenCmd)
 	prCmd.AddCommand(prMergeCmd)
+	prMergeCmd.Flags().BoolP("merge", "m", true, "merge the commits with the base branch")
+	prMergeCmd.Flags().BoolP("rebase", "r", false, "Rebase the commits onto the base branch")
+	prMergeCmd.Flags().BoolP("squash", "s", false, "Squash the commits into one commit and merge it into the base branch")
 
 	prCmd.AddCommand(prListCmd)
 	prListCmd.Flags().IntP("limit", "L", 30, "Maximum number of items to fetch")
@@ -58,7 +61,7 @@ var prStatusCmd = &cobra.Command{
 	RunE:  prStatus,
 }
 var prViewCmd = &cobra.Command{
-	Use:   "view [{<number> | <url> | <branch>}]",
+	Use:   "view <{<number> | <url> | <branch>}>",
 	Short: "View a pull request",
 	Long: `Display the title, body, and other information about a pull request.
 
@@ -82,7 +85,7 @@ var prReopenCmd = &cobra.Command{
 }
 
 var prMergeCmd = &cobra.Command{
-	Use:   "merge <number | url>",
+	Use:   "merge <number | url> [--rebase | --merge | --squash]",
 	Short: "Merge a pull request",
 	Args:  cobra.ExactArgs(1),
 	RunE:  prMerge,
@@ -450,12 +453,32 @@ func prMerge(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	err = api.PullRequestMerge(apiClient, baseRepo, pr)
+	rebase, err := cmd.Flags().GetBool("rebase")
+	if err != nil {
+		return err
+	}
+	squash, err := cmd.Flags().GetBool("squash")
+	if err != nil {
+		return err
+	}
+
+	var output string
+	if rebase {
+		output = fmt.Sprintf("%s Rebased pull request #%d\n", utils.Green("✔"), pr.Number)
+		err = api.PullRequestRebase(apiClient, baseRepo, pr)
+	} else if squash {
+		output = fmt.Sprintf("%s Squashed and merged pull request #%d\n", utils.Green("✔"), pr.Number)
+		err = api.PullRequestSquash(apiClient, baseRepo, pr)
+	} else {
+		output = fmt.Sprintf("%s Merged pull request #%d\n", utils.Green("✔"), pr.Number)
+		err = api.PullRequestMerge(apiClient, baseRepo, pr)
+	}
+
 	if err != nil {
 		return fmt.Errorf("API call failed: %w", err)
 	}
 
-	fmt.Fprintf(colorableErr(cmd), "%s Merged pull request #%d\n", utils.Green("✔"), pr.Number)
+	fmt.Fprintf(colorableErr(cmd), output)
 
 	return nil
 }
diff --git a/command/pr_test.go b/command/pr_test.go
index 86d2829d2..c4d4cfb07 100644
--- a/command/pr_test.go
+++ b/command/pr_test.go
@@ -975,3 +975,43 @@ func TestPrMerge(t *testing.T) {
 		t.Fatalf("output did not match regexp /%s/\n> output\n%q\n", r, output.Stderr())
 	}
 }
+
+func TestPrMerge_rebase(t *testing.T) {
+	initWithStubs(
+		stubResponse{200, bytes.NewBufferString(`{ "data": { "repository": {
+			"pullRequest": { "number": 2, "closed": false, "state": "OPEN"}
+		} } }`)},
+		stubResponse{200, bytes.NewBufferString(`{"id": "THE-ID"}`)},
+	)
+
+	output, err := RunCommand("pr merge 2 --rebase")
+	if err != nil {
+		t.Fatalf("error running command `pr merge`: %v", err)
+	}
+
+	r := regexp.MustCompile(`Rebased pull request #2`)
+
+	if !r.MatchString(output.Stderr()) {
+		t.Fatalf("output did not match regexp /%s/\n> output\n%q\n", r, output.Stderr())
+	}
+}
+
+func TestPrMerge_squash(t *testing.T) {
+	initWithStubs(
+		stubResponse{200, bytes.NewBufferString(`{ "data": { "repository": {
+			"pullRequest": { "number": 3, "closed": false, "state": "OPEN"}
+		} } }`)},
+		stubResponse{200, bytes.NewBufferString(`{"id": "THE-ID"}`)},
+	)
+
+	output, err := RunCommand("pr merge 3 --squash")
+	if err != nil {
+		t.Fatalf("error running command `pr merge`: %v", err)
+	}
+
+	r := regexp.MustCompile(`Squashed and merged pull request #3`)
+
+	if !r.MatchString(output.Stderr()) {
+		t.Fatalf("output did not match regexp /%s/\n> output\n%q\n", r, output.Stderr())
+	}
+}