mock: Conditionally toggle mock Call

mock/mock.go

@@ -208,10 +208,49 @@ func (c *Call) On(methodName string, arguments ...interface{}) *Call {
 	return c.Parent.On(methodName, arguments...)
 }
 
+// unsetConfiguration stores Unset method call configuration.
+type unsetConfiguration struct {
+	enabled bool
+}
+
+// UnsetOption will allow to configure Unset method invocation.
+type UnsetOption func(*unsetConfiguration)
+
+// WithUnsetToggle will configure unset value. If enabled == false the Unset
+// function will not be called. Default: true - the Unset will be called.
+// This option may be used to conditionally enable or disable mock calls.
+func WithUnsetToggle(enabled bool) UnsetOption {
+	return func(uc *unsetConfiguration) {
+		uc.enabled = enabled
+	}
+}
+
+// WithUnsetEnabled is shorthand for WithUnsetToggle(true).
+func WithUnsetEnabled() UnsetOption {
+	return func(uc *unsetConfiguration) {
+		uc.enabled = true
+	}
+}
+
+// WithUnsetEnabled is shorthand for WithUnsetToggle(false).
+func WithUnsetDisabled() UnsetOption {
+	return func(uc *unsetConfiguration) {
+		uc.enabled = false
+	}
+}
+
 // Unset removes a mock handler from being called.
 //
 //	test.On("func", mock.Anything).Unset()
-func (c *Call) Unset() *Call {
+func (c *Call) Unset(options ...UnsetOption) *Call {
+	configuration := &unsetConfiguration{enabled: true}
+	for _, option := range options {
+		option(configuration)
+	}
+	if !configuration.enabled {
+		return c
+	}
+
 	var unlockOnce sync.Once
 
 	for _, arg := range c.Arguments {
@@ -252,6 +291,15 @@ func (c *Call) Unset() *Call {
 	return c
 }
 
+// Toggle will conditionally enable or disable mock call via flag.
+// By default Call is toggled.
+func (c *Call) Toggle(t bool) *Call {
+	if t {
+		return c
+	}
+	return c.Unset()
+}
+
 // NotBefore indicates that the mock should only be called after the referenced
 // calls have been called as expected. The referenced calls may be from the
 // same mock instance and/or other mock instances.

mock/mock_test.go

@@ -666,6 +666,44 @@ func Test_Mock_Unset_WithFuncPanics(t *testing.T) {
 	})
 }
 
+func Test_Mock_Unset_WithUnsetToggle_Option(t *testing.T) {
+	// make a test impl object
+	var mockedService = new(TestExampleImplementation)
+	mockedService.On("TheExampleMethod", 1).
+		Unset(WithUnsetToggle(true))
+	assert.Len(t, mockedService.ExpectedCalls, 0)
+
+	mockedService.On("TestExampleMethod", 1).
+		Unset(WithUnsetToggle(false))
+	assert.Len(t, mockedService.ExpectedCalls, 1)
+}
+
+func Test_Mock_Unset_WithUnsetEnabled_Option(t *testing.T) {
+	// make a test impl object
+	var mockedService = new(TestExampleImplementation)
+	mockedService.On("TheExampleMethod", 1).
+		Unset(WithUnsetEnabled())
+
+	assert.Len(t, mockedService.ExpectedCalls, 0)
+}
+
+func Test_Mock_Unset_WithUnsetDisabled_Option(t *testing.T) {
+	// make a test impl object
+	var mockedService = new(TestExampleImplementation)
+	mockedService.On("TheExampleMethod", 1).
+		Unset(WithUnsetDisabled())
+	assert.Len(t, mockedService.ExpectedCalls, 1)
+}
+
+func Test_Mock_Toggle(t *testing.T) {
+	var mockedService = new(TestExampleImplementation)
+	mockedService.On("TheExampleMethod", 1).Toggle(true)
+	assert.Len(t, mockedService.ExpectedCalls, 1)
+
+	mockedService.On("TheExampleMethod", 1).Toggle(false)
+	assert.Len(t, mockedService.ExpectedCalls, 0)
+}
+
 func Test_Mock_Return(t *testing.T) {
 
 	// make a test impl object