support ctx as input params, error as output for MockClient.Call

client/mock/mock.go

@@ -84,10 +84,24 @@ func (m *MockClient) Call(ctx context.Context, req client.Request, rsp interface
 		response := r.Response
 		if t := reflect.TypeOf(r.Response); t.Kind() == reflect.Func {
 			var request []reflect.Value
-			if t.NumIn() == 1 {
+			switch t.NumIn() {
+			case 1:
+				// one input params: (req)
 				request = append(request, reflect.ValueOf(req.Body()))
+			case 2:
+				// two input params: (ctx, req)
+				request = append(request, reflect.ValueOf(ctx), reflect.ValueOf(req.Body()))
+			}
+
+			responseValue := reflect.ValueOf(r.Response).Call(request)
+			response = responseValue[0].Interface()
+			if len(responseValue) == 2 {
+				// make it possible to return error in response function
+				respErr, ok := responseValue[1].Interface().(error)
+				if ok && respErr != nil {
+					return respErr
+				}
 			}
-			response = reflect.ValueOf(r.Response).Call(request)[0].Interface()
 		}
 
 		v.Set(reflect.ValueOf(response))

client/mock/mock_test.go

@@ -2,6 +2,7 @@ package mock
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/micro/go-micro/errors"
@@ -24,6 +25,12 @@ func TestClient(t *testing.T) {
 			}
 			return "wrong"
 		}},
+		{Endpoint: "Foo.FuncWithRequestContextAndResponse", Response: func(ctx context.Context, req interface{}) string {
+			return "something"
+		}},
+		{Endpoint: "Foo.FuncWithRequestContextAndResponseError", Response: func(ctx context.Context, req interface{}) (string, error) {
+			return "something", fmt.Errorf("mock error")
+		}},
 	}
 
 	c := NewClient(Response("go.mock", response))
@@ -35,8 +42,10 @@ func TestClient(t *testing.T) {
 		err := c.Call(context.TODO(), req, &rsp)
 
 		if err != r.Error {
+			if r.Endpoint != "Foo.FuncWithRequestContextAndResponseError" {
 				t.Fatalf("Expecter error %v got %v", r.Error, err)
 			}
+		}
 
 		t.Log(rsp)
 		if r.Endpoint == "Foo.FuncWithReqBody" {