How to mock __import__ built-in function

So lets say we need to write unit test for piece of code that uses  __import__  built-in function:

def task_we_want_to_test(task_name: str, task_class: str, env: dict) -> Task:
    module = __import__(
        'my_project.scripts.{}'.format(task_name),
        fromlist=[task_class, ])
    klass = getattr(module, task_class)
    return klass(env)

 

You can test this code like this:

def test_get_task(self):
    from my_project.scripts import task_we_want_to_test

    MyScript = mock.MagicMock()
    _import = mock.MagicMock()
    _import.MyScript = MyScript
    with mock.patch.dict(
        'sys.modules',
        {'my_project.scripts.my_script': _import},
    ):
        klass = task_we_want_to_test('my_script', self.env)
        MyScript.assert_called_with(self.env)
        self.assertEqual(klass, MyScript(self.env))

 

Leave a Reply

Your email address will not be published. Required fields are marked *