From 4fcf13dac465e319c4565ff311bca8a5614dad6b Mon Sep 17 00:00:00 2001
From: Corey Johnson <probablycorey@gmail.com>
Date: Wed, 8 Jan 2020 11:44:27 -0800
Subject: [PATCH] Make it work with PRs

---
 api/queries_pr.go | 27 +++++++++++++++++++++------
 command/pr.go     | 10 ++++++----
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/api/queries_pr.go b/api/queries_pr.go
index 4b8f5fcb7..22d0ec3fd 100644
--- a/api/queries_pr.go
+++ b/api/queries_pr.go
@@ -6,11 +6,16 @@ import (
 )
 
 type PullRequestsPayload struct {
-	ViewerCreated   []PullRequest
-	ReviewRequested []PullRequest
+	ViewerCreated   PullRequestAndTotalCount
+	ReviewRequested PullRequestAndTotalCount
 	CurrentPR       *PullRequest
 }
 
+type PullRequestAndTotalCount struct {
+	TotalCount   int
+	PullRequests []PullRequest
+}
+
 type PullRequest struct {
 	Number      int
 	Title       string
@@ -123,7 +128,8 @@ type Repo interface {
 
 func PullRequests(client *Client, ghRepo Repo, currentPRNumber int, currentPRHeadRef, currentUsername string) (*PullRequestsPayload, error) {
 	type edges struct {
-		Edges []struct {
+		TotalCount int
+		Edges      []struct {
 			Node PullRequest
 		}
 	}
@@ -177,6 +183,7 @@ func PullRequests(client *Client, ghRepo Repo, currentPRNumber int, currentPRHea
 	query($owner: String!, $repo: String!, $headRefName: String!, $viewerQuery: String!, $reviewerQuery: String!, $per_page: Int = 10) {
 		repository(owner: $owner, name: $repo) {
 			pullRequests(headRefName: $headRefName, states: OPEN, first: $per_page) {
+				totalCount
 				edges {
 					node {
 						...prWithReviews
@@ -198,6 +205,7 @@ func PullRequests(client *Client, ghRepo Repo, currentPRNumber int, currentPRHea
 
 	query := fragments + queryPrefix + `
       viewerCreated: search(query: $viewerQuery, type: ISSUE, first: $per_page) {
+       totalCount: issueCount
         edges {
           node {
             ...prWithReviews
@@ -205,6 +213,7 @@ func PullRequests(client *Client, ghRepo Repo, currentPRNumber int, currentPRHea
         }
       }
       reviewRequested: search(query: $reviewerQuery, type: ISSUE, first: $per_page) {
+        totalCount: issueCount
         edges {
           node {
             ...pr
@@ -260,9 +269,15 @@ func PullRequests(client *Client, ghRepo Repo, currentPRNumber int, currentPRHea
 	}
 
 	payload := PullRequestsPayload{
-		viewerCreated,
-		reviewRequested,
-		currentPR,
+		ViewerCreated: PullRequestAndTotalCount{
+			PullRequests: viewerCreated,
+			TotalCount:   resp.ViewerCreated.TotalCount,
+		},
+		ReviewRequested: PullRequestAndTotalCount{
+			PullRequests: reviewRequested,
+			TotalCount:   resp.ReviewRequested.TotalCount,
+		},
+		CurrentPR: currentPR,
 	}
 
 	return &payload, nil
diff --git a/command/pr.go b/command/pr.go
index d63f3379a..cb20bf804 100644
--- a/command/pr.go
+++ b/command/pr.go
@@ -97,6 +97,8 @@ func prStatus(cmd *cobra.Command, args []string) error {
 
 	out := colorableOut(cmd)
 
+	fmt.Printf("🌭 %+v\n", prPayload)
+
 	printHeader(out, "Current branch")
 	if prPayload.CurrentPR != nil {
 		printPrs(out, *prPayload.CurrentPR)
@@ -107,16 +109,16 @@ func prStatus(cmd *cobra.Command, args []string) error {
 	fmt.Fprintln(out)
 
 	printHeader(out, "Created by you")
-	if len(prPayload.ViewerCreated) > 0 {
-		printPrs(out, prPayload.ViewerCreated...)
+	if prPayload.ViewerCreated.TotalCount > 0 {
+		printPrs(out, prPayload.ViewerCreated.PullRequests...)
 	} else {
 		printMessage(out, "  You have no open pull requests")
 	}
 	fmt.Fprintln(out)
 
 	printHeader(out, "Requesting a code review from you")
-	if len(prPayload.ReviewRequested) > 0 {
-		printPrs(out, prPayload.ReviewRequested...)
+	if prPayload.ReviewRequested.TotalCount > 0 {
+		printPrs(out, prPayload.ReviewRequested.PullRequests...)
 	} else {
 		printMessage(out, "  You have no pull requests to review")
 	}