Merge pull request #4 from pepabo/with_close

Add `with_close` option to Add `Close()` method

README.md

@@ -19,3 +19,4 @@ Generate mock client ([client_mock.go](example/gen/go/myapp/client_mock.go)) tha
 | --- | --- | --- |
 | `package` | string | Specify package name (ex `--go-client_opt=package=xxxx` ) |
 | `same_package` | bool | Make the package the same as the package generated by proto-gen-go (ex `--go-client_opt=same_package` ) |
+| `with_close` | bool | Add `Close()` method for closing *grpc.ClientConn from client |

example/buf.gen.yaml

@@ -15,7 +15,7 @@ plugins:
     opt: paths=source_relative
   - plugin: go-client
     out: gen/go
-    opt: paths=source_relative,same_package
+    opt: paths=source_relative,same_package,with_close
   - plugin: go-client-mock
     out: gen/go
-    opt: paths=source_relative,same_package
+    opt: paths=source_relative,same_package,with_close

example/example_test.go

@@ -42,8 +42,13 @@ func TestClientUser(t *testing.T) {
 func TestClientProject(t *testing.T) {
 	ctx := context.Background()
 	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
+	t.Cleanup(func() {
+		ctrl.Finish()
+	})
 	client := myapp.NewMockClient(ctrl)
+	t.Cleanup(func() {
+		client.Close()
+	})
 	req := &myapp.ListProjectsRequest{
 		Page: 1,
 	}
@@ -57,6 +62,8 @@ func TestClientProject(t *testing.T) {
 	lc.EXPECT().Recv().Return(res[1], nil)
 	lc.EXPECT().Recv().Return(nil, io.EOF)
 	client.EXPECT().ProjectService().ListProjects(ctx, req).Return(lc, nil)
+	client.EXPECT().Close().Return(nil)
+
 	rc, err := client.ProjectService().ListProjects(ctx, req)
 	if err != nil {
 		t.Fatal(err)

gen/gen.go

@@ -15,11 +15,13 @@ import (
 const (
 	optPackage     = "package"
 	optSamePackage = "same_package"
+	optWithClose   = "with_close"
 )
 
 type Generator struct {
 	genp        *protogen.Plugin
 	samePackage bool
+	withClose   bool
 	packageName string
 }
 
@@ -56,10 +58,13 @@ func (gen *Generator) Generate() error {
 	}
 	g.P("")
 	g.P(`import (`)
+	if gen.withClose {
+		g.P(`"reflect"`)
+		g.P("")
+	}
 	if !gen.samePackage {
 		g.P(fmt.Sprintf("%s %s", tmppf.GoPackageName, tmppf.GoImportPath.String()))
 	}
-	g.P("")
 	g.P(`"github.com/golang/mock/gomock"`)
 	g.P(`)`)
 	g.P("")
@@ -155,6 +160,21 @@ func (gen *Generator) Generate() error {
 			}
 		}
 	}
+	if gen.withClose {
+		g.P("func (m *MockClient) Close() error {")
+		g.P("m.ctrl.T.Helper()")
+		g.P(`ret := m.ctrl.Call(m, "Close")`)
+		g.P(`ret0, _ := ret[0].(error)`)
+		g.P(`return ret0`)
+		g.P("}")
+		g.P("")
+
+		g.P("func (mr *MockClientMockRecorder) Close() *gomock.Call {")
+		g.P("mr.mock.ctrl.T.Helper()")
+		g.P(`return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockClient)(nil).Close))`)
+		g.P("}")
+		g.P("")
+	}
 
 	return nil
 }
@@ -166,6 +186,8 @@ func (gen *Generator) parseOpts() error {
 		switch {
 		case o == optSamePackage:
 			gen.samePackage = true
+		case o == optWithClose:
+			gen.withClose = true
 		case strings.HasPrefix(o, fmt.Sprintf("%s=", optPackage)):
 			gen.packageName = strings.TrimPrefix(o, fmt.Sprintf("%s=", optPackage))
 		}