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()