mock_args.py

import mock

from unittest import TestCase


def expensive_function(x):
    """The function to be mocked. depending on how it is called, it may return
    different data from the database, which can take a long time, which is why we
    want to mock it."""
    if x == 'greeting':
        # do some crazy database access stuff that takes forever
        return {'hello': 'crazy_db_result'}
    if x == 'greeting':
        # do some crazy database access stuff that takes forever
        return {'goodbye': 'crazy_db_result'}


def mock_args():
    """This function is the function we are testing. It makes multiple calls to
    expensive_function, and does something different based on what is returned."""
    if expensive_function('greeting').get('hello') == 'hello':
        print expensive_function('salutation').get('goodbye')

class MockArgs(TestCase):

    def test_mock_args(self):
        with mock.patch('mock_args.expensive_function') as fake:
            # How do I set the response based on the arguments that get passed in?
            #fake('greeting') = {'hello': 'hello'}
            #fake('salutation') = {'goodbye': 'goodbye'}
            mock_args()
            fake.assert_called_with('salutation')
        with mock.patch('mock_args.expensive_function') as fake:
            #fake('greeting') = {'hello': 'nope'}
            #fake('salutation') = {'goodbye': 'goodbye'}
            mock_args()
            fake.assert_called_once()