Implement OverridableController to allow for mock overrides

Author: styeung

gomock/callset.go

@@ -27,6 +27,8 @@ type callSet struct {
 	expected map[callSetKey][]*Call
 	// Calls that have been exhausted.
 	exhausted map[callSetKey][]*Call
+	// when set to true,
+	allowOverride bool
 }
 
 // callSetKey is the key in the maps in callSet
@@ -36,7 +38,18 @@ type callSetKey struct {
 }
 
 func newCallSet() *callSet {
-	return &callSet{make(map[callSetKey][]*Call), make(map[callSetKey][]*Call)}
+	return &callSet{
+		expected:  make(map[callSetKey][]*Call),
+		exhausted: make(map[callSetKey][]*Call),
+	}
+}
+
+func newOverridableCallSet() *callSet {
+	return &callSet{
+		expected:      make(map[callSetKey][]*Call),
+		exhausted:     make(map[callSetKey][]*Call),
+		allowOverride: true,
+	}
 }
 
 // Add adds a new expected call.
@@ -46,6 +59,11 @@ func (cs callSet) Add(call *Call) {
 	if call.exhausted() {
 		m = cs.exhausted
 	}
+	if cs.allowOverride {
+		delete(m, key)
+		m[key] = make([]*Call, 0)
+	}
+
 	m[key] = append(m[key], call)
 }
 

gomock/callset_test.go

@@ -42,6 +42,24 @@ func TestCallSetAdd(t *testing.T) {
 	}
 }
 
+func TestCallSetAdd_WhenOverridable_ClearsPreviousExpectedAndExhausted(t *testing.T) {
+	method := "TestMethod"
+	var receiver interface{} = "TestReceiver"
+	cs := newOverridableCallSet()
+
+	cs.Add(newCall(t, receiver, method, reflect.TypeOf(receiverType{}.Func)))
+	numExpectedCalls := len(cs.expected[callSetKey{receiver, method}])
+	if numExpectedCalls != 1 {
+		t.Fatalf("Expected 1 expected call in callset, got %d", numExpectedCalls)
+	}
+
+	cs.Add(newCall(t, receiver, method, reflect.TypeOf(receiverType{}.Func)))
+	newNumExpectedCalls := len(cs.expected[callSetKey{receiver, method}])
+	if newNumExpectedCalls != 1 {
+		t.Fatalf("Expected 1 expected call in callset, got %d", newNumExpectedCalls)
+	}
+}
+
 func TestCallSetRemove(t *testing.T) {
 	method := "TestMethod"
 	var receiver interface{} = "TestReceiver"

gomock/controller.go

@@ -105,6 +105,26 @@ func NewController(t TestReporter) *Controller {
 	return ctrl
 }
 
+// NewOverridableController returns a new Controller that allows for overridable call expectations
+func NewOverridableController(t TestReporter) *Controller {
+	h, ok := t.(TestHelper)
+	if !ok {
+		h = &nopTestHelper{t}
+	}
+	ctrl := &Controller{
+		T:             h,
+		expectedCalls: newOverridableCallSet(),
+	}
+	if c, ok := isCleanuper(ctrl.T); ok {
+		c.Cleanup(func() {
+			ctrl.T.Helper()
+			ctrl.finish(true, nil)
+		})
+	}
+
+	return ctrl
+}
+
 type cancelReporter struct {
 	t      TestHelper
 	cancel func()

gomock/example_test.go

@@ -48,3 +48,21 @@ func ExampleCall_DoAndReturn_captureArguments() {
 	fmt.Printf("%s %s", r, s)
 	// Output: I'm sleepy foo
 }
+
+func ExampleCall_OverridableController_HappyPath() {
+	t := &testing.T{} // provided by test
+	ctrl := gomock.NewOverridableController(t)
+	mockIndex := NewMockFoo(ctrl)
+	var s string
+
+	mockIndex.EXPECT().Bar(gomock.AssignableToTypeOf(s)).DoAndReturn(
+		func(arg string) interface{} {
+			s = arg
+			return "I'm sleepy"
+		},
+	)
+
+	r := mockIndex.Bar("foo")
+	fmt.Printf("%s %s", r, s)
+	// Output: I'm sleepy foo
+}

gomock/overridable_controller_test.go

@@ -0,0 +1,34 @@
+package gomock_test
+
+import (
+	"testing"
+
+	"github.com/golang/mock/gomock"
+)
+
+func TestEcho_NoOverride(t *testing.T) {
+	ctrl := gomock.NewOverridableController(t)
+	mockIndex := NewMockFoo(ctrl)
+
+	mockIndex.EXPECT().Bar(gomock.Any()).Return("foo")
+	res := mockIndex.Bar("input")
+
+	if res != "foo" {
+		t.Fatalf("expected response to equal 'foo', got %s", res)
+	}
+}
+
+func TestEcho_WithOverride_BaseCase(t *testing.T) {
+	ctrl := gomock.NewOverridableController(t)
+	mockIndex := NewMockFoo(ctrl)
+
+	// initial expectation set
+	mockIndex.EXPECT().Bar(gomock.Any()).Return("foo")
+	// override
+	mockIndex.EXPECT().Bar(gomock.Any()).Return("bar")
+	res := mockIndex.Bar("input")
+
+	if res != "bar" {
+		t.Fatalf("expected response to equal 'bar', got %s", res)
+	}
+}