feat: add --max-items flag for paginated/streaming endpoints

Author: stainless-app[bot]

internal/mocktest/mocktest.go

@@ -15,6 +15,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -84,46 +85,12 @@ func TestRunMockTestWithPipeAndFlags(t *testing.T, pipeData []byte, args ...stri
 
 	t.Logf("Testing command: go run ./cmd/beeper-desktop-cli %s", strings.Join(args[2:], " "))
 
-	cliCmd := exec.Command("go", args...)
-	cliCmd.Stdin = bytes.NewReader(pipeData)
-
-	// Pipe the CLI tool's output into `head` so it doesn't hang when simulating
-	// paginated or streamed endpoints. 100 lines of output should be enough to
-	// test that the API endpoint worked, or report back a meaningful amount of
-	// data if something went wrong.
-	headCmd := exec.Command("head", "-n", "100")
-	pipe, err := cliCmd.StdoutPipe()
-	require.NoError(t, err, "Failed to create pipe for CLI command")
-	headCmd.Stdin = pipe
-
-	// Capture `head` output and CLI command stderr outputs:
-	var output strings.Builder
-	headCmd.Stdout = &output
-	headCmd.Stderr = &output
-	cliCmd.Stderr = &output
-
-	// First start `head`, so it's ready for data to come in:
-	err = headCmd.Start()
-	require.NoError(t, err, "Failed to start `head` command")
-
-	// Next start the CLI command so it can pipe data to `head` without
-	// buffering any data in advance:
-	err = cliCmd.Start()
-	require.NoError(t, err, "Failed to start CLI command")
-
-	// Ensure that the stdout pipe is closed as soon as `head` exits, to let the
-	// CLI tool know that no more output is needed and it can stop streaming
-	// test data for streaming/paginated endpoints. This needs to happen before
-	// calling `cliCmd.Wait()`, otherwise there will be a deadlock.
-	err = headCmd.Wait()
-	pipe.Close()
-	require.NoError(t, err, "`head` command finished with an error")
-
-	// Finally, wait for the CLI tool to finish up:
-	err = cliCmd.Wait()
-	require.NoError(t, err, "CLI command failed\n%s", output.String())
-
-	t.Logf("Test passed successfully\nOutput:\n%s", output.String())
+	cmd := exec.Command("go", args...)
+	cmd.Stdin = bytes.NewReader(pipeData)
+	output, err := cmd.CombinedOutput()
+	assert.NoError(t, err, "Test failed\nError: %v\nOutput: %s", err, output)
+
+	t.Logf("Test passed successfully\nOutput:\n%s", string(output))
 }
 
 func TestFile(t *testing.T, contents string) string {

pkg/cmd/accountcontact.go

@@ -46,6 +46,10 @@ var accountsContactsList = cli.Command{
 			Usage:     "Optional search query for blended contact lookup.",
 			QueryPath: "query",
 		},
+		&requestflag.Flag[int64]{
+			Name:  "max-items",
+			Usage: "The maximum number of items to return (use -1 for unlimited).",
+		},
 	},
 	Action:          handleAccountsContactsList,
 	HideHelpCommand: true,
@@ -119,7 +123,11 @@ func handleAccountsContactsList(ctx context.Context, cmd *cli.Command) error {
 			params,
 			options...,
 		)
-		return ShowJSONIterator(os.Stdout, "accounts:contacts list", iter, format, transform)
+		maxItems := int64(-1)
+		if cmd.IsSet("max-items") {
+			maxItems = cmd.Value("max-items").(int64)
+		}
+		return ShowJSONIterator(os.Stdout, "accounts:contacts list", iter, format, transform, maxItems)
 	}
 }
 

pkg/cmd/accountcontact_test.go

@@ -13,6 +13,7 @@ func TestAccountsContactsList(t *testing.T) {
 		mocktest.TestRunMockTestWithFlags(
 			t, "accounts:contacts", "list",
 			"--access-token", "string",
+			"--max-items", "10",
 			"--account-id", "accountID",
 			"--cursor", "1725489123456|c29tZUltc2dQYWdl",
 			"--direction", "before",

pkg/cmd/chat.go

@@ -136,6 +136,10 @@ var chatsList = cli.Command{
 			Usage:     "Pagination direction used with 'cursor': 'before' fetches older results, 'after' fetches newer results. Defaults to 'before' when only 'cursor' is provided.",
 			QueryPath: "direction",
 		},
+		&requestflag.Flag[int64]{
+			Name:  "max-items",
+			Usage: "The maximum number of items to return (use -1 for unlimited).",
+		},
 	},
 	Action:          handleChatsList,
 	HideHelpCommand: true,
@@ -231,6 +235,10 @@ var chatsSearch = cli.Command{
 			Usage:     "Set to true to only retrieve chats that have unread messages",
 			QueryPath: "unreadOnly",
 		},
+		&requestflag.Flag[int64]{
+			Name:  "max-items",
+			Usage: "The maximum number of items to return (use -1 for unlimited).",
+		},
 	},
 	Action:          handleChatsSearch,
 	HideHelpCommand: true,
@@ -346,7 +354,11 @@ func handleChatsList(ctx context.Context, cmd *cli.Command) error {
 		return ShowJSON(os.Stdout, "chats list", obj, format, transform)
 	} else {
 		iter := client.Chats.ListAutoPaging(ctx, params, options...)
-		return ShowJSONIterator(os.Stdout, "chats list", iter, format, transform)
+		maxItems := int64(-1)
+		if cmd.IsSet("max-items") {
+			maxItems = cmd.Value("max-items").(int64)
+		}
+		return ShowJSONIterator(os.Stdout, "chats list", iter, format, transform, maxItems)
 	}
 }
 
@@ -416,6 +428,10 @@ func handleChatsSearch(ctx context.Context, cmd *cli.Command) error {
 		return ShowJSON(os.Stdout, "chats search", obj, format, transform)
 	} else {
 		iter := client.Chats.SearchAutoPaging(ctx, params, options...)
-		return ShowJSONIterator(os.Stdout, "chats search", iter, format, transform)
+		maxItems := int64(-1)
+		if cmd.IsSet("max-items") {
+			maxItems = cmd.Value("max-items").(int64)
+		}
+		return ShowJSONIterator(os.Stdout, "chats search", iter, format, transform, maxItems)
 	}
 }

pkg/cmd/chat_test.go

@@ -88,6 +88,7 @@ func TestChatsList(t *testing.T) {
 		mocktest.TestRunMockTestWithFlags(
 			t, "chats", "list",
 			"--access-token", "string",
+			"--max-items", "10",
 			"--account-id", "local-whatsapp_ba_EvYDBBsZbRQAy3UOSWqG0LuTVkc",
 			"--account-id", "local-instagram_ba_eRfQMmnSNy_p7Ih7HL7RduRpKFU",
 			"--cursor", "1725489123456|c29tZUltc2dQYWdl",
@@ -122,6 +123,7 @@ func TestChatsSearch(t *testing.T) {
 		mocktest.TestRunMockTestWithFlags(
 			t, "chats", "search",
 			"--access-token", "string",
+			"--max-items", "10",
 			"--account-id", "local-whatsapp_ba_EvYDBBsZbRQAy3UOSWqG0LuTVkc",
 			"--account-id", "local-telegram_ba_QFrb5lrLPhO3OT5MFBeTWv0x4BI",
 			"--cursor", "1725489123456|c29tZUltc2dQYWdl",

pkg/cmd/cmdutil.go

@@ -350,7 +350,7 @@ type HasRawJSON interface {
 
 // For an iterator over different value types, display its values to the user in
 // different formats.
-func ShowJSONIterator[T any](stdout *os.File, title string, iter jsonview.Iterator[T], format string, transform string) error {
+func ShowJSONIterator[T any](stdout *os.File, title string, iter jsonview.Iterator[T], format string, transform string, itemsToDisplay int64) error {
 	if format == "explore" {
 		return jsonview.ExploreJSONStream(title, iter)
 	}
@@ -366,6 +366,9 @@ func ShowJSONIterator[T any](stdout *os.File, title string, iter jsonview.Iterat
 	output := []byte{}
 	numberOfNewlines := 0
 	for iter.Next() {
+		if itemsToDisplay == 0 {
+			break
+		}
 		item := iter.Current()
 		var obj gjson.Result
 		if hasRaw, ok := any(item).(HasRawJSON); ok {
@@ -383,6 +386,7 @@ func ShowJSONIterator[T any](stdout *os.File, title string, iter jsonview.Iterat
 		}
 
 		output = append(output, json...)
+		itemsToDisplay -= 1
 		numberOfNewlines += countTerminalLines(json, terminalWidth)
 
 		// If the output won't fit in the terminal window, stream it to a pager
@@ -409,6 +413,9 @@ func ShowJSONIterator[T any](stdout *os.File, title string, iter jsonview.Iterat
 		}
 
 		for iter.Next() {
+			if itemsToDisplay == 0 {
+				break
+			}
 			item := iter.Current()
 			var obj gjson.Result
 			if hasRaw, ok := any(item).(HasRawJSON); ok {
@@ -423,6 +430,7 @@ func ShowJSONIterator[T any](stdout *os.File, title string, iter jsonview.Iterat
 			if err := ShowJSON(pager, title, obj, format, transform); err != nil {
 				return err
 			}
+			itemsToDisplay -= 1
 		}
 		return iter.Err()
 	})

pkg/cmd/message.go

@@ -60,6 +60,10 @@ var messagesList = cli.Command{
 			Usage:     "Pagination direction used with 'cursor': 'before' fetches older results, 'after' fetches newer results. Defaults to 'before' when only 'cursor' is provided.",
 			QueryPath: "direction",
 		},
+		&requestflag.Flag[int64]{
+			Name:  "max-items",
+			Usage: "The maximum number of items to return (use -1 for unlimited).",
+		},
 	},
 	Action:          handleMessagesList,
 	HideHelpCommand: true,
@@ -138,6 +142,10 @@ var messagesSearch = cli.Command{
 			Usage:     "Filter by sender: 'me' (messages sent by the authenticated user), 'others' (messages sent by others), or a specific user ID string (user.id).",
 			QueryPath: "sender",
 		},
+		&requestflag.Flag[int64]{
+			Name:  "max-items",
+			Usage: "The maximum number of items to return (use -1 for unlimited).",
+		},
 	},
 	Action:          handleMessagesSearch,
 	HideHelpCommand: true,
@@ -297,7 +305,11 @@ func handleMessagesList(ctx context.Context, cmd *cli.Command) error {
 			params,
 			options...,
 		)
-		return ShowJSONIterator(os.Stdout, "messages list", iter, format, transform)
+		maxItems := int64(-1)
+		if cmd.IsSet("max-items") {
+			maxItems = cmd.Value("max-items").(int64)
+		}
+		return ShowJSONIterator(os.Stdout, "messages list", iter, format, transform, maxItems)
 	}
 }
 
@@ -335,7 +347,11 @@ func handleMessagesSearch(ctx context.Context, cmd *cli.Command) error {
 		return ShowJSON(os.Stdout, "messages search", obj, format, transform)
 	} else {
 		iter := client.Messages.SearchAutoPaging(ctx, params, options...)
-		return ShowJSONIterator(os.Stdout, "messages search", iter, format, transform)
+		maxItems := int64(-1)
+		if cmd.IsSet("max-items") {
+			maxItems = cmd.Value("max-items").(int64)
+		}
+		return ShowJSONIterator(os.Stdout, "messages search", iter, format, transform, maxItems)
 	}
 }
 

pkg/cmd/message_test.go

@@ -37,6 +37,7 @@ func TestMessagesList(t *testing.T) {
 		mocktest.TestRunMockTestWithFlags(
 			t, "messages", "list",
 			"--access-token", "string",
+			"--max-items", "10",
 			"--chat-id", "!NCdzlIaMjZUmvmvyHU:beeper.com",
 			"--cursor", "1725489123456|c29tZUltc2dQYWdl",
 			"--direction", "before",
@@ -49,6 +50,7 @@ func TestMessagesSearch(t *testing.T) {
 		mocktest.TestRunMockTestWithFlags(
 			t, "messages", "search",
 			"--access-token", "string",
+			"--max-items", "10",
 			"--account-id", "local-whatsapp_ba_EvYDBBsZbRQAy3UOSWqG0LuTVkc",
 			"--account-id", "local-instagram_ba_eRfQMmnSNy_p7Ih7HL7RduRpKFU",
 			"--chat-id", "!NCdzlIaMjZUmvmvyHU:beeper.com",