Rework mock service and tests to accommadate mockgen

codegen/templates/module_mock_initializer.tmpl

@@ -5,16 +5,21 @@
 package module
 
 import (
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	zanzibar "github.com/uber/zanzibar/runtime"
+
 	{{range $classType, $moduleInstances := $instance.RecursiveDependencies -}}
 	{{range $idx, $moduleInstance := $moduleInstances -}}
+	{{if eq $classType $leafClass -}}
+	{{$moduleInstance.PackageInfo.ImportPackageAlias}} "{{$moduleInstance.PackageInfo.ImportPackagePath}}/mock_client"
+	{{else -}}
 	{{$moduleInstance.PackageInfo.ImportPackageAlias}} "{{$moduleInstance.PackageInfo.ImportPackagePath}}"
-	{{if not (eq $classType $leafClass) -}}
 	{{$moduleInstance.PackageInfo.ModulePackageAlias}} "{{$moduleInstance.PackageInfo.ModulePackagePath}}"
 	{{end -}}
 	{{end -}}
 	{{end}}
-
-	zanzibar "github.com/uber/zanzibar/runtime"
 )
 
 {{$moduleInstances := (index $instance.RecursiveDependencies $leafClass) -}}
@@ -29,6 +34,7 @@ type {{$mockDeps}} struct {
 // for the {{$instance.InstanceName}} {{$instance.ClassName}} with leaf nodes being mocks
 func InitializeDependenciesMock(
 	g *zanzibar.Gateway,
+	ctrl *gomock.Controller,
 ) (*DependenciesTree, *Dependencies, *{{$mockDeps}}) {
 	tree := &DependenciesTree{}
 
@@ -45,7 +51,7 @@ func InitializeDependenciesMock(
 	{{camel $mockDeps}} := &{{$mockDeps}}{
 		{{- range $idx, $dependency := $moduleInstances}}
 		{{- $pkgInfo := $dependency.PackageInfo }}
-		{{$pkgInfo.QualifiedInstanceName}}: &{{$pkgInfo.ImportPackageAlias}}.Mock{{title $className}}{},
+		{{$pkgInfo.QualifiedInstanceName}}: {{$pkgInfo.ImportPackageAlias}}.NewMock{{title $className}}(ctrl),
 		{{- end }}
 	}
 	{{- $initializedDeps := printf "initialized%sDependencies" (title $className) }}

codegen/templates/service_mock.tmpl

@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"path/filepath"
 	"runtime"
+	"testing"
 	"time"
 
 
@@ -43,13 +44,14 @@ type MockService interface {
 type mockService struct {
 	started        bool
 	server         *zanzibar.Gateway
+	ctrl           *gomock.Controller
 	{{camel $mockType}}    *module.{{$mockType}}
 	httpClient     *http.Client
 	tChannelClient zanzibar.TChannelCaller
 }
 
 // MustCreateTestService creates a new MockService, panics if it fails doing so.
-func MustCreateTestService() MockService {
+func MustCreateTestService(t *testing.T) MockService {
 	_, file, _, _ := runtime.Caller(0)
 	currentDir := zanzibar.GetDirnameFromRuntimeCaller(file)
 	testConfigPath := filepath.Join(currentDir, "../../../config/test.json")
@@ -60,7 +62,8 @@ func MustCreateTestService() MockService {
 		panic(err)
 	}
 
-	_, dependencies, mockNodes := module.InitializeDependenciesMock(server)
+	ctrl := gomock.NewController(t)
+	_, dependencies, mockNodes := module.InitializeDependenciesMock(server, ctrl)
 	registerErr := registerDeps(server, dependencies)
 	if registerErr != nil {
 		panic(registerErr)
@@ -91,6 +94,7 @@ func MustCreateTestService() MockService {
 
 	return &mockService{
 		server:         		server,
+		ctrl:                   ctrl,
 		{{camel $mockType}}:    mockNodes,
 		httpClient:     		httpClient,
 		tChannelClient: 		tchannelClient,
@@ -107,6 +111,7 @@ func (m *mockService) Start() {
 
 // Stop stops the mock server
 func (m *mockService) Stop() {
+	m.ctrl.Finish()
 	m.server.Close()
 	m.started = false
 }

examples/example-gateway/endpoints/contacts/save_contacts_test.go

@@ -4,16 +4,16 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
 
 	clientContacts "github.com/uber/zanzibar/examples/example-gateway/build/gen-code/clients/contacts/contacts"
 	endpointContacts "github.com/uber/zanzibar/examples/example-gateway/build/gen-code/endpoints/contacts/contacts"
 	"github.com/uber/zanzibar/examples/example-gateway/build/services/example-gateway"
 )
 
 func TestSaveContactsCall(t *testing.T) {
-	ms := examplegatewayServiceGenerated.MustCreateTestService()
+	ms := examplegatewayServiceGenerated.MustCreateTestService(t)
 	ms.Start()
 	defer ms.Stop()
 
@@ -27,7 +27,7 @@ func TestSaveContactsCall(t *testing.T) {
 	}
 	clientResponse := &clientContacts.SaveContactsResponse{}
 
-	ms.MockClientNodes().Contacts.On("SaveContacts", mock.Anything, mock.Anything, clientRequest).
+	ms.MockClientNodes().Contacts.EXPECT().SaveContacts(gomock.Any(), gomock.Any(), clientRequest).
 		Return(clientResponse, nil, nil)
 
 	res, err := ms.MakeHTTPRequest(

examples/example-gateway/endpoints/tchannel/baz/baz_call_test.go

@@ -4,15 +4,15 @@ import (
 	"context"
 	"testing"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
 
 	"github.com/uber/zanzibar/examples/example-gateway/build/gen-code/endpoints/tchannel/baz/baz"
 	ms "github.com/uber/zanzibar/examples/example-gateway/build/services/example-gateway"
 )
 
 func TestBazCall(t *testing.T) {
-	ms := ms.MustCreateTestService()
+	ms := ms.MustCreateTestService(t)
 	ms.Start()
 	defer ms.Stop()
 
@@ -34,7 +34,7 @@ func TestBazCall(t *testing.T) {
 
 	ctx := context.Background()
 	var result baz.SimpleService_Call_Result
-	ms.MockClientNodes().Baz.On("Call", mock.Anything, reqHeaders, mock.Anything).
+	ms.MockClientNodes().Baz.EXPECT().Call(gomock.Any(), reqHeaders, gomock.Any()).
 		Return(map[string]string{"some-res-header": "something"}, nil)
 
 	success, resHeaders, err := ms.MakeTChannelRequest(