Repository: golang/mock
Pull request #111 (2017-09-20 23:15:59)

gomock.WithContext: controller with a context
gomock.WithContext: controller with a context which is cancelled on Fatals.

testing.T.FailNow (and Fatal, ...) work only when called from the same
goroutine as the test. This makes using gomock with e.g. RPC server
mocks quite a challenge as the tests are likely to deadlock instead of
failing properly.

This adds a new constructor returning an additional context, which is
cancelled on any fatal failure in mock.

This provides a reasonable workaround for #110 and similar issues.


Patch:
diff --git a/gomock/controller.go b/gomock/controller.go
index d06dfbb..14069e5 100644
--- a/gomock/controller.go
+++ b/gomock/controller.go
@@ -56,6 +56,7 @@
 package gomock
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"runtime"
@@ -85,6 +86,24 @@ func NewController(t TestReporter) *Controller {
 	}
 }
 
+type cancelReporter struct {
+	t      TestReporter
+	cancel func()
+}
+
+func (r *cancelReporter) Errorf(format string, args ...interface{}) { r.t.Errorf(format, args...) }
+func (r *cancelReporter) Fatalf(format string, args ...interface{}) {
+	defer r.cancel()
+	r.t.Fatalf(format, args...)
+}
+
+// WithContext returns a new Controller and a Context, which is cancelled on any
+// fatal failure.
+func WithContext(ctx context.Context, t TestReporter) (*Controller, context.Context) {
+	ctx, cancel := context.WithCancel(ctx)
+	return NewController(&cancelReporter{t, cancel}), ctx
+}
+
 func (ctrl *Controller) RecordCall(receiver interface{}, method string, args ...interface{}) *Call {
 	recv := reflect.ValueOf(receiver)
 	for i := 0; i < recv.Type().NumMethod(); i++ {
diff --git a/sample/concurrent/concurrent.go b/sample/concurrent/concurrent.go
new file mode 100644
index 0000000..346cee4
--- /dev/null
+++ b/sample/concurrent/concurrent.go
@@ -0,0 +1,8 @@
+//go:generate mockgen -destination mock/concurrent_mock.go github.com/golang/mock/sample/concurrent Math
+
+// Package concurrent demonstrates how to use gomock with goroutines.
+package concurrent
+
+type Math interface {
+	Sum(a, b int) int
+}
diff --git a/sample/concurrent/mock/concurrent_mock.go b/sample/concurrent/mock/concurrent_mock.go
new file mode 100644
index 0000000..9200563
--- /dev/null
+++ b/sample/concurrent/mock/concurrent_mock.go
@@ -0,0 +1,45 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/golang/mock/sample/concurrent (interfaces: Math)
+
+// Package mock_concurrent is a generated GoMock package.
+package mock_concurrent
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+)
+
+// MockMath is a mock of Math interface
+type MockMath struct {
+	ctrl     *gomock.Controller
+	recorder *MockMathMockRecorder
+}
+
+// MockMathMockRecorder is the mock recorder for MockMath
+type MockMathMockRecorder struct {
+	mock *MockMath
+}
+
+// NewMockMath creates a new mock instance
+func NewMockMath(ctrl *gomock.Controller) *MockMath {
+	mock := &MockMath{ctrl: ctrl}
+	mock.recorder = &MockMathMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockMath) EXPECT() *MockMathMockRecorder {
+	return m.recorder
+}
+
+// Sum mocks base method
+func (m *MockMath) Sum(arg0, arg1 int) int {
+	ret := m.ctrl.Call(m, "Sum", arg0, arg1)
+	ret0, _ := ret[0].(int)
+	return ret0
+}
+
+// Sum indicates an expected call of Sum
+func (mr *MockMathMockRecorder) Sum(arg0, arg1 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Sum", reflect.TypeOf((*MockMath)(nil).Sum), arg0, arg1)
+}


Test patch:
diff --git a/sample/concurrent/concurrent_test.go b/sample/concurrent/concurrent_test.go
new file mode 100644
index 0000000..2cc0780
--- /dev/null
+++ b/sample/concurrent/concurrent_test.go
@@ -0,0 +1,45 @@
+package concurrent
+
+import (
+	"context"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	mock "github.com/golang/mock/sample/concurrent/mock"
+)
+
+func call(ctx context.Context, m Math) (int, error) {
+	result := make(chan int)
+	go func() {
+		result <- m.Sum(1, 2)
+		close(result)
+	}()
+	select {
+	case r := <-result:
+		return r, nil
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	}
+}
+
+// testConcurrentFails is expected to fail (and is disabled). It
+// demonstrates how to use gomock.WithContext to interrupt the test
+// from a different goroutine.
+func testConcurrentFails(t *testing.T) {
+	ctrl, ctx := gomock.WithContext(context.Background(), t)
+	defer ctrl.Finish()
+	m := mock.NewMockMath(ctrl)
+	if _, err := call(ctx, m); err != nil {
+		t.Error("call failed:", err)
+	}
+}
+
+func TestConcurrentWorks(t *testing.T) {
+	ctrl, ctx := gomock.WithContext(context.Background(), t)
+	defer ctrl.Finish()
+	m := mock.NewMockMath(ctrl)
+	m.EXPECT().Sum(1, 2).Return(3)
+	if _, err := call(ctx, m); err != nil {
+		t.Error("call failed:", err)
+	}
+}