Repository: googleapis/gapic-generator-go
Pull request #880 (2022-02-01 18:52:55)

feat(gengapic): diregapic lro examples use wrapper
This updates the generated GoDoc examples for DIREGAPIC LRO RPCs to use the operation wrapper returned by the client to `Wait` for operation completion.

This also cleans up/refactors some of the internal helper code/tests.


Patch:
diff --git a/internal/gengapic/custom_operation.go b/internal/gengapic/custom_operation.go
index 011b297..938ae27 100644
--- a/internal/gengapic/custom_operation.go
+++ b/internal/gengapic/custom_operation.go
@@ -36,8 +36,9 @@ func (g *generator) isCustomOp(m *descriptor.MethodDescriptorProto, info *httpIn
 	return g.opts.diregapic && // Generator in DIREGAPIC mode.
 		g.aux.customOp != nil && // API Defines a custom operation.
 		m.GetOutputType() == g.customOpProtoName() && // Method returns the custom operation.
-		info.verb != "get" && // Method is not a GET (polling methods).
-		m.GetName() != "Wait" // Method is not a Wait (uses POST).
+		m.GetName() != "Wait" && // Method is not a Wait (uses POST).
+		info != nil && // Must have google.api.http.
+		info.verb != "get" // Method is not a GET (polling methods).
 }
 
 // customOpProtoName builds the fully-qualified proto name for the custom
diff --git a/internal/gengapic/example.go b/internal/gengapic/example.go
index 167bc20..1eddd7a 100644
--- a/internal/gengapic/example.go
+++ b/internal/gengapic/example.go
@@ -105,6 +105,8 @@ func (g *generator) exampleMethodBody(pkgName, servName string, m *descriptor.Me
 		return err
 	}
 
+	httpInfo := getHTTPInfo(m)
+
 	g.imports[inSpec] = true
 	// Pick the first transport for simplicity. We don't need examples
 	// of each method for both transports when they have the same surface.
@@ -129,7 +131,7 @@ func (g *generator) exampleMethodBody(pkgName, servName string, m *descriptor.Me
 	}
 	if pf != nil {
 		g.examplePagingCall(m)
-	} else if g.isLRO(m) {
+	} else if g.isLRO(m) || g.isCustomOp(m, httpInfo) {
 		g.exampleLROCall(m)
 	} else if *m.OutputType == emptyType {
 		g.exampleEmptyCall(m)
@@ -149,7 +151,7 @@ func (g *generator) exampleLROCall(m *descriptor.MethodDescriptorProto) {
 	// if response_type is google.protobuf.Empty, don't generate a "resp" var
 	eLRO := proto.GetExtension(m.Options, longrunning.E_OperationInfo)
 	opInfo := eLRO.(*longrunning.OperationInfo)
-	if opInfo.GetResponseType() == emptyValue {
+	if opInfo.GetResponseType() == emptyValue || opInfo == nil {
 		// no new variables when this is used
 		// therefore don't attempt to delcare it
 		retVars = "err ="
diff --git a/internal/gengapic/gengapic.go b/internal/gengapic/gengapic.go
index 4a7c2d2..532e399 100644
--- a/internal/gengapic/gengapic.go
+++ b/internal/gengapic/gengapic.go
@@ -484,10 +484,7 @@ func (g *generator) returnType(m *descriptor.MethodDescriptorProto) (string, err
 	if err != nil {
 		return "", err
 	}
-	info, err := getHTTPInfo(m)
-	if err != nil {
-		return "", err
-	}
+	info := getHTTPInfo(m)
 
 	// Regular return type.
 	retTyp := fmt.Sprintf("*%s.%s", outSpec.Name, outType.GetName())
diff --git a/internal/gengapic/genrest.go b/internal/gengapic/genrest.go
index 44b45ff..905402a 100644
--- a/internal/gengapic/genrest.go
+++ b/internal/gengapic/genrest.go
@@ -194,8 +194,8 @@ type httpInfo struct {
 
 func (g *generator) pathParams(m *descriptor.MethodDescriptorProto) map[string]*descriptor.FieldDescriptorProto {
 	pathParams := map[string]*descriptor.FieldDescriptorProto{}
-	info, err := getHTTPInfo(m)
-	if info == nil || err != nil {
+	info := getHTTPInfo(m)
+	if info == nil {
 		return pathParams
 	}
 
@@ -218,8 +218,8 @@ func (g *generator) pathParams(m *descriptor.MethodDescriptorProto) map[string]*
 
 func (g *generator) queryParams(m *descriptor.MethodDescriptorProto) map[string]*descriptor.FieldDescriptorProto {
 	queryParams := map[string]*descriptor.FieldDescriptorProto{}
-	info, err := getHTTPInfo(m)
-	if info == nil || err != nil {
+	info := getHTTPInfo(m)
+	if info == nil {
 		return queryParams
 	}
 	if info.body == "*" {
@@ -395,10 +395,7 @@ func (g *generator) generateQueryString(m *descriptor.MethodDescriptorProto) {
 }
 
 func (g *generator) generateURLString(m *descriptor.MethodDescriptorProto) error {
-	info, err := getHTTPInfo(m)
-	if err != nil {
-		return err
-	}
+	info := getHTTPInfo(m)
 	if info == nil {
 		return errors.E(nil, "method has no http info: %s", m.GetName())
 	}
@@ -427,9 +424,9 @@ func (g *generator) generateURLString(m *descriptor.MethodDescriptorProto) error
 	return nil
 }
 
-func getHTTPInfo(m *descriptor.MethodDescriptorProto) (*httpInfo, error) {
+func getHTTPInfo(m *descriptor.MethodDescriptorProto) *httpInfo {
 	if m == nil || m.GetOptions() == nil {
-		return nil, nil
+		return nil
 	}
 
 	eHTTP := proto.GetExtension(m.GetOptions(), annotations.E_Http)
@@ -455,7 +452,7 @@ func getHTTPInfo(m *descriptor.MethodDescriptorProto) (*httpInfo, error) {
 		info.url = httpRule.GetDelete()
 	}
 
-	return &info, nil
+	return &info
 }
 
 // genRESTMethod generates a single method from a client. m must be a method declared in serv.
@@ -561,7 +558,7 @@ func (g *generator) pagingRESTCall(servName string, m *descriptor.MethodDescript
 	if err != nil {
 		return err
 	}
-	info, err := getHTTPInfo(m)
+	info := getHTTPInfo(m)
 	if err != nil {
 		return err
 	}
@@ -692,10 +689,7 @@ func (g *generator) lroRESTCall(servName string, m *descriptor.MethodDescriptorP
 }
 
 func (g *generator) emptyUnaryRESTCall(servName string, m *descriptor.MethodDescriptorProto) error {
-	info, err := getHTTPInfo(m)
-	if err != nil {
-		return err
-	}
+	info := getHTTPInfo(m)
 	if info == nil {
 		return errors.E(nil, "method has no http info: %s", m.GetName())
 	}
@@ -771,10 +765,7 @@ func (g *generator) emptyUnaryRESTCall(servName string, m *descriptor.MethodDesc
 }
 
 func (g *generator) unaryRESTCall(servName string, m *descriptor.MethodDescriptorProto) error {
-	info, err := getHTTPInfo(m)
-	if err != nil {
-		return err
-	}
+	info := getHTTPInfo(m)
 	if info == nil {
 		return errors.E(nil, "method has no http info: %s", m.GetName())
 	}


Test patch:
diff --git a/internal/gengapic/example_test.go b/internal/gengapic/example_test.go
index 6b1efcb..07984e5 100644
--- a/internal/gengapic/example_test.go
+++ b/internal/gengapic/example_test.go
@@ -22,6 +22,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
 	"github.com/googleapis/gapic-generator-go/internal/txtdiff"
+	"google.golang.org/genproto/googleapis/api/annotations"
 	"google.golang.org/genproto/googleapis/api/serviceconfig"
 	"google.golang.org/genproto/googleapis/longrunning"
 	"google.golang.org/protobuf/proto"
@@ -30,13 +31,14 @@ import (
 
 func TestExample(t *testing.T) {
 	var g generator
-	g.opts = &options{}
 	g.imports = map[pbinfo.ImportSpec]bool{}
-	g.mixins = mixins{
+	g.aux = &auxTypes{}
+	mix := mixins{
 		"google.longrunning.Operations":   operationsMethods(),
 		"google.cloud.location.Locations": locationMethods(),
 		"google.iam.v1.IAMPolicy":         iamPolicyMethods(),
 	}
+	g.mixins = mix
 	g.serviceConfig = &serviceconfig.Service{
 		Apis: []*apipb.Api{
 			{Name: "foo.bar.Baz"},
@@ -84,10 +86,15 @@ func TestExample(t *testing.T) {
 		},
 	}
 
+	cop := &descriptor.DescriptorProto{
+		Name: proto.String("Operation"),
+	}
+
 	file := &descriptor.FileDescriptorProto{
 		Options: &descriptor.FileOptions{
 			GoPackage: proto.String("mypackage"),
 		},
+		Package: proto.String("my.pkg"),
 	}
 
 	emptyLRO := &longrunning.OperationInfo{
@@ -102,11 +109,18 @@ func TestExample(t *testing.T) {
 	respLROOpts := &descriptor.MethodOptions{}
 	proto.SetExtension(respLROOpts, longrunning.E_OperationInfo, respLRO)
 
+	customOpOpts := &descriptor.MethodOptions{}
+	proto.SetExtension(customOpOpts, annotations.E_Http, &annotations.HttpRule{
+		Pattern: &annotations.HttpRule_Post{
+			Post: "/customOp",
+		},
+	})
+
 	commonTypes(&g)
 	for _, typ := range []*descriptor.DescriptorProto{
-		inputType, outputType, pageInputType, pageOutputType,
+		inputType, outputType, pageInputType, pageOutputType, cop,
 	} {
-		g.descInfo.Type[".my.pkg."+*typ.Name] = typ
+		g.descInfo.Type[".my.pkg."+typ.GetName()] = typ
 		g.descInfo.ParentFile[typ] = file
 	}
 
@@ -127,6 +141,7 @@ func TestExample(t *testing.T) {
 				Name:       proto.String("GetBigThing"),
 				InputType:  proto.String(".my.pkg.InputType"),
 				OutputType: proto.String(".google.longrunning.Operation"),
+				Options:    respLROOpts,
 			},
 			{
 				Name:       proto.String("GetManyThings"),
@@ -152,17 +167,26 @@ func TestExample(t *testing.T) {
 				OutputType: proto.String(".google.longrunning.Operation"),
 				Options:    respLROOpts,
 			},
+			{
+				Name:       proto.String("CustomOp"),
+				InputType:  proto.String(".my.pkg.InputType"),
+				OutputType: proto.String(".my.pkg.Operation"),
+				Options:    customOpOpts,
+			},
 		},
 	}
 	for _, tst := range []struct {
-		tstName, pkgName string
-		transports       []transport
-		imports          map[pbinfo.ImportSpec]bool
+		tstName string
+		imports map[pbinfo.ImportSpec]bool
+		options options
+		op      *customOp
 	}{
 		{
-			tstName:    "empty_example",
-			pkgName:    "Foo",
-			transports: []transport{grpc, rest},
+			tstName: "empty_example",
+			options: options{
+				pkgName:    "Foo",
+				transports: []transport{grpc, rest},
+			},
 			imports: map[pbinfo.ImportSpec]bool{
 				{Path: "context"}:                        true,
 				{Path: "google.golang.org/api/iterator"}: true,
@@ -174,9 +198,11 @@ func TestExample(t *testing.T) {
 			},
 		},
 		{
-			tstName:    "empty_example_grpc",
-			pkgName:    "Foo",
-			transports: []transport{grpc},
+			tstName: "empty_example_grpc",
+			options: options{
+				pkgName:    "Foo",
+				transports: []transport{grpc},
+			},
 			imports: map[pbinfo.ImportSpec]bool{
 				{Path: "context"}:                        true,
 				{Path: "google.golang.org/api/iterator"}: true,
@@ -188,9 +214,11 @@ func TestExample(t *testing.T) {
 			},
 		},
 		{
-			tstName:    "foo_example",
-			pkgName:    "Bar",
-			transports: []transport{grpc, rest},
+			tstName: "foo_example",
+			options: options{
+				pkgName:    "Bar",
+				transports: []transport{grpc, rest},
+			},
 			imports: map[pbinfo.ImportSpec]bool{
 				{Path: "context"}:                        true,
 				{Path: "google.golang.org/api/iterator"}: true,
@@ -202,9 +230,11 @@ func TestExample(t *testing.T) {
 			},
 		},
 		{
-			tstName:    "foo_example_rest",
-			pkgName:    "Bar",
-			transports: []transport{rest},
+			tstName: "foo_example_rest",
+			options: options{
+				pkgName:    "Bar",
+				transports: []transport{rest},
+			},
 			imports: map[pbinfo.ImportSpec]bool{
 				{Path: "context"}:                        true,
 				{Path: "google.golang.org/api/iterator"}: true,
@@ -215,10 +245,29 @@ func TestExample(t *testing.T) {
 				{Name: "mypackagepb", Path: "mypackage"}:                                           true,
 			},
 		},
+		{
+			tstName: "custom_op_example",
+			options: options{
+				pkgName:    "Bar",
+				transports: []transport{rest},
+				diregapic:  true,
+			},
+			imports: map[pbinfo.ImportSpec]bool{
+				{Path: "context"}:                        true,
+				{Path: "google.golang.org/api/iterator"}: true,
+				{Path: "io"}:                             true,
+				{Name: "mypackagepb", Path: "mypackage"}: true,
+			},
+			op: &customOp{message: cop},
+		},
 	} {
 		g.reset()
-		g.opts.pkgName = tst.pkgName
-		g.opts.transports = tst.transports
+		g.opts = &tst.options
+		g.mixins = mix
+		if tst.options.diregapic {
+			g.mixins = nil
+		}
+		g.aux.customOp = tst.op
 		g.genExampleFile(serv)
 		if diff := cmp.Diff(g.imports, tst.imports); diff != "" {
 			t.Errorf("TestExample(%s): imports got(-),want(+):\n%s", tst.tstName, diff)
diff --git a/internal/gengapic/testdata/custom_op_example.want b/internal/gengapic/testdata/custom_op_example.want
new file mode 100644
index 0000000..0f6985e
--- /dev/null
+++ b/internal/gengapic/testdata/custom_op_example.want
@@ -0,0 +1,207 @@
+func ExampleNewFooRESTClient() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	// TODO: Use client.
+	_ = c
+}
+
+func ExampleFooClient_GetEmptyThing() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	err = c.GetEmptyThing(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+}
+
+func ExampleFooClient_GetOneThing() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	resp, err := c.GetOneThing(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	// TODO: Use resp.
+	_ = resp
+}
+
+func ExampleFooClient_GetBigThing() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	op, err := c.GetBigThing(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+
+	resp, err := op.Wait(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	// TODO: Use resp.
+	_ = resp
+}
+
+func ExampleFooClient_GetManyThings() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.PageInputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#PageInputType.
+	}
+	it := c.GetManyThings(ctx, req)
+	for {
+		resp, err := it.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			// TODO: Handle error.
+		}
+		// TODO: Use resp.
+		_ = resp
+	}
+}
+
+func ExampleFooClient_BidiThings() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+	stream, err := c.BidiThings(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	go func() {
+		reqs := []*mypackagepb.InputType{
+			// TODO: Create requests.
+		}
+		for _, req := range reqs {
+			if err := stream.Send(req); err != nil {
+				// TODO: Handle error.
+			}
+		}
+		stream.CloseSend()
+	}()
+	for {
+		resp, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			// TODO: handle error.
+		}
+		// TODO: Use resp.
+		_ = resp
+	}
+}
+
+func ExampleFooClient_EmptyLRO() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	op, err := c.EmptyLRO(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+
+	err = op.Wait(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+}
+
+func ExampleFooClient_RespLRO() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	op, err := c.RespLRO(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+
+	resp, err := op.Wait(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	// TODO: Use resp.
+	_ = resp
+}
+
+func ExampleFooClient_CustomOp() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	op, err := c.CustomOp(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+
+	err = op.Wait(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+}
+
diff --git a/internal/gengapic/testdata/empty_example.want b/internal/gengapic/testdata/empty_example.want
index b3fe8bf..624b7e9 100644
--- a/internal/gengapic/testdata/empty_example.want
+++ b/internal/gengapic/testdata/empty_example.want
@@ -194,6 +194,26 @@ func ExampleClient_RespLRO() {
 	_ = resp
 }
 
+func ExampleClient_CustomOp() {
+	ctx := context.Background()
+	c, err := Foo.NewClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	resp, err := c.CustomOp(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	// TODO: Use resp.
+	_ = resp
+}
+
 func ExampleClient_ListLocations() {
 	ctx := context.Background()
 	c, err := Foo.NewClient(ctx)
diff --git a/internal/gengapic/testdata/empty_example_grpc.want b/internal/gengapic/testdata/empty_example_grpc.want
index 34db3af..6fcd6c4 100644
--- a/internal/gengapic/testdata/empty_example_grpc.want
+++ b/internal/gengapic/testdata/empty_example_grpc.want
@@ -182,6 +182,26 @@ func ExampleClient_RespLRO() {
 	_ = resp
 }
 
+func ExampleClient_CustomOp() {
+	ctx := context.Background()
+	c, err := Foo.NewClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	resp, err := c.CustomOp(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	// TODO: Use resp.
+	_ = resp
+}
+
 func ExampleClient_ListLocations() {
 	ctx := context.Background()
 	c, err := Foo.NewClient(ctx)
diff --git a/internal/gengapic/testdata/foo_example.want b/internal/gengapic/testdata/foo_example.want
index 3207a36..8c79e5d 100644
--- a/internal/gengapic/testdata/foo_example.want
+++ b/internal/gengapic/testdata/foo_example.want
@@ -194,6 +194,26 @@ func ExampleFooClient_RespLRO() {
 	_ = resp
 }
 
+func ExampleFooClient_CustomOp() {
+	ctx := context.Background()
+	c, err := Bar.NewFooClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	resp, err := c.CustomOp(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	// TODO: Use resp.
+	_ = resp
+}
+
 func ExampleFooClient_ListLocations() {
 	ctx := context.Background()
 	c, err := Bar.NewFooClient(ctx)
diff --git a/internal/gengapic/testdata/foo_example_rest.want b/internal/gengapic/testdata/foo_example_rest.want
index 54c5d9d..64bad87 100644
--- a/internal/gengapic/testdata/foo_example_rest.want
+++ b/internal/gengapic/testdata/foo_example_rest.want
@@ -182,6 +182,26 @@ func ExampleFooClient_RespLRO() {
 	_ = resp
 }
 
+func ExampleFooClient_CustomOp() {
+	ctx := context.Background()
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.InputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#InputType.
+	}
+	resp, err := c.CustomOp(ctx, req)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	// TODO: Use resp.
+	_ = resp
+}
+
 func ExampleFooClient_ListLocations() {
 	ctx := context.Background()
 	c, err := Bar.NewFooRESTClient(ctx)