fix(gencli): guard iterator.Done check for LROs with paging fields (#1758)

Author: hongalex

internal/gencli/cmd_file.go

@@ -364,7 +364,7 @@ var {{$methodCmdVar}} = &cobra.Command{
 		printMessage(result){{ end }}
 		{{ end }}
 		{{ end }}
-		{{ if .Paged }}
+		{{ if and .Paged (not .IsLRO) }}
 		if err == iterator.Done {
 		    return nil
 		}

internal/gencli/cmd_file_test.go

@@ -137,6 +137,41 @@ func TestCommandFile(t *testing.T) {
 		IsLRO: true,
 	}
 
+	// LRO and Paged
+	lroPagedCmd := &Command{
+		Service:           "Todo",
+		Method:            "LROPagedTodo",
+		MethodCmd:         "lro-paged-todo",
+		ShortDesc:         "lro paged todo",
+		LongDesc:          "lro paged todo description",
+		InputMessage:      "todopb.LROPagedTodoRequest",
+		InputMessageVar:   "LROPagedTodoInput",
+		OutputMessageType: ".google.longrunning.Operation",
+		Imports: map[string]*pbinfo.ImportSpec{
+			"todopb": &pbinfo.ImportSpec{Name: "todopb", Path: "github.com/googleapis/todo/generated"},
+		},
+		Flags: []*Flag{
+			&Flag{
+				Name:      "page_size",
+				Type:      descriptorpb.FieldDescriptorProto_TYPE_INT32,
+				FieldName: "PageSize",
+				VarName:   "LROPagedTodoInput",
+				Usage:     "size of page to list",
+				Optional:  true,
+			},
+			&Flag{
+				Name:      "page_token",
+				FieldName: "PageToken",
+				VarName:   "LROPagedTodoInput",
+				Type:      descriptorpb.FieldDescriptorProto_TYPE_STRING,
+				Usage:     "the token from a previous page listed",
+				Optional:  true,
+			},
+		},
+		IsLRO: true,
+		Paged: true,
+	}
+
 	// client streaming
 	copyTodosCmd := &Command{
 		Service:           "Todo",
@@ -224,6 +259,14 @@ func TestCommandFile(t *testing.T) {
 			name:       "start-todo",
 			goldenPath: filepath.Join("testdata", "start-todo.want"),
 		},
+		{
+			g: &gcli{
+				format: true,
+			},
+			cmd:        lroPagedCmd,
+			name:       "lro-paged-todo",
+			goldenPath: filepath.Join("testdata", "lro-paged.want"),
+		},
 		{
 			g: &gcli{
 				format: true,

internal/gencli/testdata/lro-paged.want

@@ -0,0 +1,147 @@
+// Code generated. DO NOT EDIT.
+
+package main
+
+import (
+	"github.com/spf13/cobra"
+
+	todopb "github.com/googleapis/todo/generated"
+)
+
+var LROPagedTodoInput todopb.LROPagedTodoRequest
+
+var LROPagedTodoFromFile string
+
+var LROPagedTodoFollow bool
+
+var LROPagedTodoPollOperation string
+
+var lROPagedTodoInputPageSize int32
+
+var lROPagedTodoInputPageToken string
+
+func init() {
+	TodoServiceCmd.AddCommand(LROPagedTodoCmd)
+
+	LROPagedTodoCmd.Flags().Int32Var(&lROPagedTodoInputPageSize, "page_size", 10, "Default is 10. size of page to list")
+
+	LROPagedTodoCmd.Flags().StringVar(&lROPagedTodoInputPageToken, "page_token", "", "the token from a previous page listed")
+
+	LROPagedTodoCmd.Flags().StringVar(&LROPagedTodoFromFile, "from_file", "", "Absolute path to JSON file containing request payload")
+
+	LROPagedTodoCmd.Flags().BoolVar(&LROPagedTodoFollow, "follow", false, "Block until the long running operation completes")
+
+	TodoServiceCmd.AddCommand(LROPagedTodoPollCmd)
+
+	LROPagedTodoPollCmd.Flags().BoolVar(&LROPagedTodoFollow, "follow", false, "Block until the long running operation completes")
+
+	LROPagedTodoPollCmd.Flags().StringVar(&LROPagedTodoPollOperation, "operation", "", "Required. Operation name to poll for")
+
+	LROPagedTodoPollCmd.MarkFlagRequired("operation")
+
+}
+
+var LROPagedTodoCmd = &cobra.Command{
+	Use:   "lro-paged-todo",
+	Short: "lro paged todo",
+	Long:  "lro paged todo description",
+	PreRun: func(cmd *cobra.Command, args []string) {
+
+		if LROPagedTodoFromFile == "" {
+
+		}
+
+	},
+	RunE: func(cmd *cobra.Command, args []string) (err error) {
+
+		in := os.Stdin
+		if LROPagedTodoFromFile != "" {
+			in, err = os.Open(LROPagedTodoFromFile)
+			if err != nil {
+				return err
+			}
+			defer in.Close()
+
+			err = jsonpb.Unmarshal(in, &LROPagedTodoInput)
+			if err != nil {
+				return err
+			}
+
+		}
+
+		if Verbose {
+			printVerboseInput("Todo", "LROPagedTodo", &LROPagedTodoInput)
+		}
+		resp, err := TodoClient.LROPagedTodo(ctx, &LROPagedTodoInput)
+		if err != nil {
+			return err
+		}
+
+		if !LROPagedTodoFollow {
+			var s interface{}
+			s = resp.Name()
+
+			if OutputJSON {
+				d := make(map[string]string)
+				d["operation"] = resp.Name()
+				s = d
+			}
+
+			printMessage(s)
+			return err
+		}
+
+		result, err := resp.Wait(ctx)
+		if err != nil {
+			return err
+		}
+
+		if Verbose {
+			fmt.Print("Output: ")
+		}
+		printMessage(result)
+
+		return err
+	},
+}
+
+var LROPagedTodoPollCmd = &cobra.Command{
+	Use:   "poll-lro-paged-todo",
+	Short: "Poll the status of a LROPagedTodoOperation by name",
+	RunE: func(cmd *cobra.Command, args []string) (err error) {
+		op := TodoClient.LROPagedTodoOperation(LROPagedTodoPollOperation)
+
+		if LROPagedTodoFollow {
+			resp, err := op.Wait(ctx)
+			if err != nil {
+				return err
+			}
+
+			if Verbose {
+				fmt.Print("Output: ")
+			}
+			printMessage(resp)
+			return err
+		}
+
+		resp, err := op.Poll(ctx)
+		if err != nil {
+			return err
+		} else if resp != nil {
+			if Verbose {
+				fmt.Print("Output: ")
+			}
+
+			printMessage(resp)
+			return
+		}
+
+		if op.Done() {
+			fmt.Println(fmt.Sprintf("Operation %s is done", op.Name()))
+		} else {
+			fmt.Println(fmt.Sprintf("Operation %s not done", op.Name()))
+		}
+
+		return err
+	},
+}