diff --git a/mock/mock.go b/mock/mock.go index d5eb1ef55..c4bc245a4 100644 --- a/mock/mock.go +++ b/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. diff --git a/mock/mock_test.go b/mock/mock_test.go index b80a8a75b..a38d9b239 100644 --- a/mock/mock_test.go +++ b/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