Don't error out for variadic functions in Do/DoReturn (#601)

Fixes: #600

Author: joshblum

gomock/call.go

@@ -108,20 +108,25 @@ func (c *Call) MaxTimes(n int) *Call {
 // DoAndReturn declares the action to run when the call is matched.
 // The return values from this function are returned by the mocked function.
 // It takes an interface{} argument to support n-arity functions.
-// The anonymous function must have the same number of input and output arguments as the mocked method.
+// The anonymous function must match the function signature mocked method.
 func (c *Call) DoAndReturn(f interface{}) *Call {
 	// TODO: Check arity and types here, rather than dying badly elsewhere.
 	v := reflect.ValueOf(f)
 
 	c.addAction(func(args []interface{}) []interface{} {
 		c.t.Helper()
-		vArgs := make([]reflect.Value, len(args))
 		ft := v.Type()
 		if c.methodType.NumIn() != ft.NumIn() {
-			c.t.Fatalf("wrong number of arguments in DoAndReturn func for %T.%v: got %d, want %d [%s]",
-				c.receiver, c.method, ft.NumIn(), c.methodType.NumIn(), c.origin)
+			if ft.IsVariadic() {
+				c.t.Fatalf("wrong number of arguments in DoAndReturn func for %T.%v The function signature must match the mocked method, a variadic function cannot be used.",
+					c.receiver, c.method)
+			} else {
+				c.t.Fatalf("wrong number of arguments in DoAndReturn func for %T.%v: got %d, want %d [%s]",
+					c.receiver, c.method, ft.NumIn(), c.methodType.NumIn(), c.origin)
+			}
 			return nil
 		}
+		vArgs := make([]reflect.Value, len(args))
 		for i := 0; i < len(args); i++ {
 			if args[i] != nil {
 				vArgs[i] = reflect.ValueOf(args[i])
@@ -144,20 +149,25 @@ func (c *Call) DoAndReturn(f interface{}) *Call {
 // return values are ignored to retain backward compatibility. To use the
 // return values call DoAndReturn.
 // It takes an interface{} argument to support n-arity functions.
-// The anonymous function must have the same number of input arguments as the mocked method.
+// The anonymous function must match the function signature mocked method.
 func (c *Call) Do(f interface{}) *Call {
 	// TODO: Check arity and types here, rather than dying badly elsewhere.
 	v := reflect.ValueOf(f)
 
 	c.addAction(func(args []interface{}) []interface{} {
 		c.t.Helper()
-		if c.methodType.NumIn() != v.Type().NumIn() {
-			c.t.Fatalf("wrong number of arguments in Do func for %T.%v: got %d, want %d [%s]",
-				c.receiver, c.method, v.Type().NumIn(), c.methodType.NumIn(), c.origin)
+		ft := v.Type()
+		if c.methodType.NumIn() != ft.NumIn() {
+			if ft.IsVariadic() {
+				c.t.Fatalf("wrong number of arguments in Do func for %T.%v The function signature must match the mocked method, a variadic function cannot be used.",
+					c.receiver, c.method)
+			} else {
+				c.t.Fatalf("wrong number of arguments in Do func for %T.%v: got %d, want %d [%s]",
+					c.receiver, c.method, ft.NumIn(), c.methodType.NumIn(), c.origin)
+			}
 			return nil
 		}
 		vArgs := make([]reflect.Value, len(args))
-		ft := v.Type()
 		for i := 0; i < len(args); i++ {
 			if args[i] != nil {
 				vArgs[i] = reflect.ValueOf(args[i])

gomock/call_test.go

@@ -510,6 +510,13 @@ func TestCall_Do_NumArgValidation(t *testing.T) {
 			args:       []interface{}{"just", "right"},
 			wantErr:    false,
 		},
+		{
+			name:       "variadic",
+			methodType: reflect.TypeOf(func(one, two string) {}),
+			doFn:       func(args ...interface{}) {},
+			args:       []interface{}{"just", "right"},
+			wantErr:    true,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -559,6 +566,13 @@ func TestCall_DoAndReturn_NumArgValidation(t *testing.T) {
 			args:       []interface{}{"just", "right"},
 			wantErr:    false,
 		},
+		{
+			name:       "variadic",
+			methodType: reflect.TypeOf(func(one, two string) {}),
+			doFn:       func(args ...interface{}) string { return "" },
+			args:       []interface{}{"just", "right"},
+			wantErr:    true,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {