From 0f2886354763fc7461c452c509ca444204e1f62d Mon Sep 17 00:00:00 2001
From: Tejaswi Agarwal <agarwalt@uber.com>
Date: Tue, 17 Sep 2019 15:30:01 -0700
Subject: [PATCH] Add option to select which mocks to generate

---
 codegen/module_system.go | 20 +++++++++++++++-----
 codegen/runner/runner.go |  2 +-
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/codegen/module_system.go b/codegen/module_system.go
index b7c09d3ac..fdf9fa00f 100644
--- a/codegen/module_system.go
+++ b/codegen/module_system.go
@@ -230,6 +230,9 @@ type ClientTestFixture struct {
 // module system (clients, endpoints, services) with a post build hook to generate client and service mocks
 func NewDefaultModuleSystemWithMockHook(
 	h *PackageHelper,
+	clientsMock bool,
+	workflowMock bool,
+	serviceMock bool,
 	hooks ...PostGenHook,
 ) (*ModuleSystem, error) {
 	t, err := NewDefaultTemplate()
@@ -237,13 +240,20 @@ func NewDefaultModuleSystemWithMockHook(
 		return nil, err
 	}
 
-	clientMockGenHook, err := ClientMockGenHook(h, t)
-	if err != nil {
-		return nil, errors.Wrap(err, "error creating client mock gen hook")
+	var clientMockGenHook, workflowMockGenHook, serviceMockGenHook PostGenHook
+	if clientsMock {
+		clientMockGenHook, err = ClientMockGenHook(h, t)
+		if err != nil {
+			return nil, errors.Wrap(err, "error creating client mock gen hook")
+		}
 	}
 
-	workflowMockGenHook := WorkflowMockGenHook(h, t)
-	serviceMockGenHook := ServiceMockGenHook(h, t)
+	if workflowMock {
+		workflowMockGenHook = WorkflowMockGenHook(h, t)
+	}
+	if serviceMock {
+		serviceMockGenHook = ServiceMockGenHook(h, t)
+	}
 
 	allHooks := append([]PostGenHook{
 		clientMockGenHook,
diff --git a/codegen/runner/runner.go b/codegen/runner/runner.go
index 8355bef90..e0cf5c6f1 100644
--- a/codegen/runner/runner.go
+++ b/codegen/runner/runner.go
@@ -142,7 +142,7 @@ func main() {
 	}
 	var moduleSystem *codegen.ModuleSystem
 	if genMock {
-		moduleSystem, err = codegen.NewDefaultModuleSystemWithMockHook(packageHelper)
+		moduleSystem, err = codegen.NewDefaultModuleSystemWithMockHook(packageHelper, true, true, true)
 	} else {
 		moduleSystem, err = codegen.NewDefaultModuleSystem(packageHelper)
 	}