gomock.WithContext: controller with a context

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++ {

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
+}

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)
+	}
+}