From 517e4b1e3fb8890cd51463d24f88705ae3aac17a Mon Sep 17 00:00:00 2001
From: Ryan Hang <rhang@uber.com>
Date: Fri, 1 Dec 2023 10:46:29 -0800
Subject: [PATCH] thriftrw-plugin-yarpc: Support go.uber.org/mock

github.com/golang/mock has been deprecated and archive and the readme
now points to https://github.com/uber/mock as its successor.

YARPC's thriftrw plugin assumes golang/mock. Add the ability to
configure the mock runtime in the thriftrw plugin generated code
to allow for golang/mock (by default to not break existing code)
and uber/mock by configuration.
---
 CHANGELOG.md                                     |  1 +
 encoding/thrift/thriftrw-plugin-yarpc/gomock.go  |  2 +-
 encoding/thrift/thriftrw-plugin-yarpc/main.go    | 16 +++++++++++++++-
 .../thrift/thriftrw-plugin-yarpc/template.go     |  1 +
 4 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e61b9cc64..a458f5de3 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/encoding/thrift/thriftrw-plugin-yarpc/gomock.go b/encoding/thrift/thriftrw-plugin-yarpc/gomock.go
index 69f584ca3..94eba8ceb 100644
--- a/encoding/thrift/thriftrw-plugin-yarpc/gomock.go
+++ b/encoding/thrift/thriftrw-plugin-yarpc/gomock.go
@@ -33,7 +33,7 @@ const gomockTemplate = `
 <$pkgname := printf "%stest" (lower .Name)>
 package <$pkgname>
 
-<$gomock := import "github.com/golang/mock/gomock">
+<$gomock := import .MockLibrary>
 
 // MockClient implements a gomock-compatible mock client for service
 // <.Name>.
diff --git a/encoding/thrift/thriftrw-plugin-yarpc/main.go b/encoding/thrift/thriftrw-plugin-yarpc/main.go
index 365287f7c..9b977ecf0 100644
--- a/encoding/thrift/thriftrw-plugin-yarpc/main.go
+++ b/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
+var (
+	_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,
diff --git a/encoding/thrift/thriftrw-plugin-yarpc/template.go b/encoding/thrift/thriftrw-plugin-yarpc/template.go
index 216d3e250..297e9a273 100644
--- a/encoding/thrift/thriftrw-plugin-yarpc/template.go
+++ b/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