thriftrw-plugin-yarpc: Support go.uber.org/mock

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ## [Unreleased]
 - grpc: returned outbound response body is no longer writable.
 - Fixed panic when error details list contains message that cannot be unmarshalled.
+- thriftrw-plugin-yarpc: Add option for configuring the mock library.
 
 ## [1.70.4] - 2023-08-31
 - logging: fix logged error in observability middleware when fields of transport.Request is in the tagsBlocklist

encoding/thrift/thriftrw-plugin-yarpc/main.go

@@ -64,12 +64,19 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"strings"
 
 	"go.uber.org/thriftrw/plugin"
 	"go.uber.org/thriftrw/plugin/api"
 )
 
+// mock libraries
+const (
+	_golangMock = "github.com/golang/mock/gomock"
+	_uberMock = "go.uber.org/mock/gomock"
+)
+
 // Command line flags
 var (
 	_context = flag.String("context-import-path",
@@ -82,7 +89,9 @@ var (
 		"go.uber.org/yarpc/encoding/thrift.OnewayHandler",
 		"Function used to wrap generic Thrift oneway function handlers into YARPC handlers")
 	_noGomock = flag.Bool("no-gomock", false,
-		"Don't generate gomock mocks for service clients")
+		"Don't generate mocks for service clients")
+	_mockLibrary = flag.String("mock-library", _golangMock,
+		fmt.Sprintf("Mock library service clients are generated with. Supported options: %q %q", _golangMock, _uberMock))
 	_noFx             = flag.Bool("no-fx", false, "Don't generate Fx module")
 	_sanitizeTChannel = flag.Bool("sanitize-tchannel", false, "Enable tchannel context sanitization")
 )
@@ -106,6 +115,10 @@ func (g g) Generate(req *api.GenerateServiceRequest) (*api.GenerateServiceRespon
 		serviceGenerators = append(serviceGenerators, gomockGenerator)
 	}
 
+	if !(*_mockLibrary == _golangMock || *_mockLibrary == _uberMock) {
+		return nil, fmt.Errorf("%q specified as mock-library. expected %q or %q", *_mockLibrary, _golangMock, _uberMock)
+	}
+
 	unaryWrapperImport, unaryWrapperFunc := splitFunctionPath(*_unaryHandlerWrapper)
 	onewayWrapperImport, onewayWrapperFunc := splitFunctionPath(*_onewayHandlerWrapper)
 
@@ -115,6 +128,7 @@ func (g g) Generate(req *api.GenerateServiceRequest) (*api.GenerateServiceRespon
 		data := serviceTemplateData{
 			Svc:                 buildSvc(serviceID, req),
 			ContextImportPath:   *_context,
+			MockLibrary:         *_mockLibrary,
 			UnaryWrapperImport:  unaryWrapperImport,
 			UnaryWrapperFunc:    unaryWrapperFunc,
 			OnewayWrapperImport: onewayWrapperImport,

encoding/thrift/thriftrw-plugin-yarpc/template.go

@@ -105,6 +105,7 @@ type serviceTemplateData struct {
 	OnewayWrapperImport string
 	OnewayWrapperFunc   string
 	SanitizeTChannel    bool
+	MockLibrary         string
 }
 
 // moduleTemplateData contains the data for code gen templates. This should be