测试 function 是否返回正确 function

Question

I have a function which dispatches other function depending on file extension in a given path.

def read_file(path: str):
    condition = function_to_check_condition(path)

    if condition == 1: return read_1(path)
    if condition == 2: return read_2(path)
    if condition == 3: return read_3(path)

def read_1(path):
    file = load_file(path)
    return file

I want to write unittest to check if it returns the correct function. I can't use assertIsInstance as it returns file . So how would I check if it returns call to read_1() first?

Answer 1

https://docs.python.org/3/library/unittest.mock.html

If you replace read_1 with a mock object named mock, you can then check mock.called or mock.assert_called_once()

Suppose your source is "so.py"

def function_to_check_condition(path):
    lookup = {
        'first': 1,
        'second': 2,
        'third': 3,
    }
    return lookup.get(path, 0)


def read_file(path: str):
    condition = function_to_check_condition(path)
    if condition == 1:
        return read_1(path)
    if condition == 2:
        return read_2(path)
    if condition == 3:
        return read_3(path)
    print ('probably should check for valid inputs, too')


def read_1(path):
    print('real method was called, but should not be by the test')
def read_2(path):
    print('real method was called, but should not be by the test')
def read_3(path):
    print('real method was called, but should not be by the test')

and suppose your test is "so_test.py" which you run with "python so_test.py"

import unittest
from unittest.mock import patch

class ReadFileTests(unittest.TestCase):

    @patch('so.read_1')
    @patch('so.read_2')
    @patch('so.read_3')
    def test_read_file(self, mock3, mock2, mock1):
        from so import read_file
        read_file('first')
        self.assertTrue(mock1.called)
        read_file('second')
        self.assertTrue(mock2.called)
        read_file('third')
        self.assertTrue(mock3.called)

        mock1.assert_called_once()
        mock2.assert_called_once()
        mock3.assert_called_once()


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