From 6a6c01a1881b6210295318bdeaa6ba97e4125b6b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Rold=C3=A1n=20Betancort?=
 <vroldanbet@authzed.com>
Date: Mon, 15 Jul 2024 09:48:15 +0100
Subject: [PATCH] PR feedback: test page size

---
 internal/commands/permission.go      |  7 ++++++-
 internal/commands/permission_test.go | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/internal/commands/permission.go b/internal/commands/permission.go
index cf1a006..5e6b4fe 100644
--- a/internal/commands/permission.go
+++ b/internal/commands/permission.go
@@ -417,6 +417,8 @@ func expandCmdFunc(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+var newLookupResourcesPageCallbackForTests func(readByPage uint)
+
 func lookupResourcesCmdFunc(cmd *cobra.Command, args []string) error {
 	objectNS := args[0]
 	relation := args[1]
@@ -493,7 +495,10 @@ func lookupResourcesCmdFunc(cmd *cobra.Command, args []string) error {
 			}
 		}
 
-		if count == 0 || pageLimit == 0 {
+		if newLookupResourcesPageCallbackForTests != nil {
+			newLookupResourcesPageCallbackForTests(count)
+		}
+		if count == 0 || pageLimit == 0 || count < uint(pageLimit) {
 			log.Trace().Interface("request", request).Uint32("page-limit", pageLimit).Uint("count", totalCount).Send()
 			break
 		}
diff --git a/internal/commands/permission_test.go b/internal/commands/permission_test.go
index 5a95bd5..a7e4198 100644
--- a/internal/commands/permission_test.go
+++ b/internal/commands/permission_test.go
@@ -138,6 +138,7 @@ func TestLookupResourcesCommand(t *testing.T) {
 	_, err = c.WriteRelationships(ctx, &v1.WriteRelationshipsRequest{Updates: updates})
 	require.NoError(t, err)
 
+	// we override this to obtain the results being printed and validate them
 	previous := console.Println
 	defer func() {
 		console.Println = previous
@@ -147,25 +148,39 @@ func TestLookupResourcesCommand(t *testing.T) {
 		count += len(values)
 	}
 
+	// use test callback to make sure pagination is correct
+	var receivedPageSizes []uint
+	newLookupResourcesPageCallbackForTests = func(readPageSize uint) {
+		receivedPageSizes = append(receivedPageSizes, readPageSize)
+	}
+	defer func() {
+		newLookupResourcesPageCallbackForTests = nil
+	}()
+
 	// test no page size, server computes returns all resources in one go
 	cmd := testLookupResourcesCommand(t, 0)
 	err = lookupResourcesCmdFunc(cmd, []string{"test/resource", "read", "test/user:1"})
 	require.NoError(t, err)
 	require.Equal(t, 10, count)
+	require.EqualValues(t, []uint{10}, receivedPageSizes)
 
 	// use page size same as number of elements
 	count = 0
+	receivedPageSizes = nil
 	cmd = testLookupResourcesCommand(t, 10)
 	err = lookupResourcesCmdFunc(cmd, []string{"test/resource", "read", "test/user:1"})
 	require.NoError(t, err)
 	require.Equal(t, 10, count)
+	require.EqualValues(t, []uint{10, 0}, receivedPageSizes)
 
 	// use odd page size
 	count = 0
+	receivedPageSizes = nil
 	cmd = testLookupResourcesCommand(t, 3)
 	err = lookupResourcesCmdFunc(cmd, []string{"test/resource", "read", "test/user:1"})
 	require.NoError(t, err)
 	require.Equal(t, 10, count)
+	require.EqualValues(t, []uint{3, 3, 3, 1}, receivedPageSizes)
 }
 
 func testLookupResourcesCommand(t *testing.T, limit uint32) *cobra.Command {