Generate custom matchers for variadic types

#90
petergtz · Jun 24, 2019 · ef093be · ef093be
1 parent 7fc7e47
commit ef093be
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
@@ -150,6 +150,9 @@ func (g *generator) generateMockFor(iface *model.Interface, mockTypeName, selfPa
 
 		addTypesFromMethodParamsTo(g.typesSet, method.In, g.packageMap)
 		addTypesFromMethodParamsTo(g.typesSet, method.Out, g.packageMap)
+		if method.Variadic != nil {
+			addTypesFromMethodParamsTo(g.typesSet, []*model.Parameter{method.Variadic}, g.packageMap)
+		}
 	}
 	g.generateMockVerifyMethods(mockTypeName)
 	g.generateVerifierType(mockTypeName)

diff --git a/mockgen/mockgen_test.go b/mockgen/mockgen_test.go
@@ -16,7 +16,7 @@ var _ = Describe("Mockgen", func() {
 			_, matcherSourceCodes := mockgen.GenerateOutput(ast, "irrelevant", "MockDisplay", "test_package", "")
 
 			Expect(matcherSourceCodes).To(SatisfyAll(
-				HaveLen(9),
+				HaveLen(10),
 				HaveKeyWithValue("http_request", SatisfyAll(
 					ContainSubstring("http \"net/http\""),
 					ContainSubstring("func AnyHttpRequest() http.Request"),
@@ -48,6 +48,9 @@ var _ = Describe("Mockgen", func() {
 				HaveKeyWithValue("send_chan_of_error", SatisfyAll(
 					ContainSubstring("func AnySendChanOfError() chan<- error"),
 				)),
+				HaveKeyWithValue("map_of_int_to_int", SatisfyAll(
+					ContainSubstring("func AnyMapOfIntToInt() map[int]int"),
+				)),
 			))
 		})
 	})

diff --git a/test_interface/display.go b/test_interface/display.go
@@ -44,4 +44,5 @@ type Display interface {
 	UseTime(t time.Time)
 	ChanParams(<-chan string, chan<- error)
 	ChanReturnValues() (<-chan string, chan<- error)
+	VariadicWithNonPrimitiveType(m ...map[int]int)
 }