diff --git a/gomock/controller.go b/gomock/controller.go index d06dfbbe..14069e58 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 00000000..346cee42 --- /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/concurrent_test.go b/sample/concurrent/concurrent_test.go new file mode 100644 index 00000000..2cc07808 --- /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) + } +} diff --git a/sample/concurrent/mock/concurrent_mock.go b/sample/concurrent/mock/concurrent_mock.go new file mode 100644 index 00000000..92005631 --- /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) +}