Python mock function with only specific argument

Question

I'm new to Python and I'm trying to mock a function only when a specific argument is passed. If other than the desired argument is passed, I'd like to call the original function instead.

In Python 2.7 I tried something like this:

from foo import config

def test_something(self):
    original_config = config # config is a Module.

    def side_effect(key):
        if key == 'expected_argument':
            return mocked_result
        else:
            return original_config.get(key)

        config.get = Mock(side_effect=side_effect)

        # actualy_test_something...

It won't work 'cause original_config is not a copy of config. It references the same module ending up in an infinite loop. I could try cloning the original config module instead but that seems to be overkill.

Is there something similar to RSpec's mocks I could use? e.g:

obj.stub(:message).with('an_expected_argument').and_return('a_mocked_result')

Any help would be appreciated. Thanks.

Answer 1

You'd need to store a reference to the unpatched function first:

def test_something(self):
    original_config_get = config.get

    def side_effect(key):
        if key == 'expected_argument':
            return mocked_result
        else:
            return original_config_get(key)

        config.get = Mock(side_effect=side_effect)

Here original_config_get references the original function before you replaced it with a Mock() object.