Mock an function to modify partial return value by special arguments on Python

Mock an function to modify partial return value by special arguments on Python

python mock一个带参数的方法，修改指定参数的返回值，大家直接看代码就能懂。

want mock code:

import requests


def get_value(arg):
    resp = requests.get('https://httpbin.org/get', params={'k': arg})
    return resp.json()['args']

def main():
    # just mock while arg == 'bad'
    print(get_value('bad'))
    print(get_value('good'))

if __name__ == '__main__':
    main()

mock code:

import unittest.mock
import mock_func


class TestMockFunc(unittest.TestCase):
    def test_main(self):
        '''style 1: using with statement and nested function'''
        print('test_main')

        # backup original function for normal call
        orig_func = mock_func.get_value

        # nested function for mock side_effect
        def fake_get_value(arg):
            if arg == 'bad':
                return {'k': 'mock_value'}
            else:
                return orig_func(arg)

        # patch function
        with unittest.mock.patch('mock_func.get_value') as mock_get_value:
            mock_get_value.side_effect = fake_get_value
            mock_func.main()

    # backup original function for normal call at class or global
    orig_func = mock_func.get_value

    # patch using decorator
    @unittest.mock.patch('mock_func.get_value')
    def test_main_other(self, mock_get_value):
        '''style 2: using with statement and nested function'''
        print('test_main_other')

        # using lambda instead of nested function
        mock_get_value.side_effect = lambda arg: 
            {'k': 'mock_value'} 
            if arg == 'bad' 
            else TestMockFunc.orig_func(arg)

        mock_func.main()

if __name__ == '__main__':
    unittest.main()