Mocking imported Python modules implementation

#python #unittest

This week, while I was developing an improvement in our team’s Ansible project, I was caught by a challenge: how to prevent any code snippet from a python module from running other than the function you are importing into the unit test target module, class or function?

The scenario encountered was during the development of a custom module for Ansible. As I was developing in a Windows environment, an Ansible dependency had problems during the import in the module ansible.module_utils.basic:

# Copyright (c), Michael DeHaan <michael.dehaan@gmail.com>, 2012-2013
# Copyright (c), Toshio Kuratomi <tkuratomi@ansible.com> 2016
# Simplified BSD License (see licenses/simplified_bsd.txt or https://opensource.org/licenses/BSD-2-Clause)

from __future__ import absolute_import, division, print_function
__metaclass__ = type

FILE_ATTRIBUTES = {
    'A': 'noatime',
    'a': 'append',
    'c': 'compressed',
    'C': 'nocow',
    'd': 'nodump',
    'D': 'dirsync',
    'e': 'extents',
    'E': 'encrypted',
    'h': 'blocksize',
    'i': 'immutable',
    'I': 'indexed',
    'j': 'journalled',
    'N': 'inline',
    's': 'zero',
    'S': 'synchronous',
    't': 'notail',
    'T': 'blockroot',
    'u': 'undelete',
    'X': 'compressedraw',
    'Z': 'compresseddirty',
}

# Ansible modules can be written in any language.
# The functions available here can be used to do many common tasks,
# to simplify development of Python modules.

import __main__
import atexit
import errno
import datetime
import grp
import fcntl
import locale
import os
import pwd
import platform
import re
import select
import shlex
import shutil
import signal
import stat
import subprocess
import sys
import tempfile
import time
import traceback
import types

The package that had the error was grp.

At the time of implementing the unit tests of my custom module, it was not relevant for me to fix this problem, since any and all functions that I would consume from the ansible.module_utils.basic module would be mocked.

For this, I spent a few hours researching a way out of this scenario. Here I come across the following solution: using the sys module, I should create a mock of the module file in question. That way, when I run my unit tests, the import in the form from … import … would not execute any code snippet from the module ansible.module_utils.basic.

Let's exemplify this solution with a very simple scenario:

We have the following Python project:

├── custom_module
│   ├── __init__.py
│   ├── module_a.py
│   ├── module_b.py
└── tests
    ├── __init__.py
    └── test_module_b.py

Function function b of module module b makes a call to function function_a of module module_a. For this, module_b needs to import the module module_a. Let's take a look at the implementation of the two modules:

# custom_module/module_a.py

def function_a():
    print("I'm the function A")

raise Exception("some generic error in module a")

# custom_module/module_b.py

from custom_module.module_a import function_a

def function_b():
    print("I'm the function B")
    function_a()

I know it's very strange for a module to simply have an exception throw, but to simulate the “defective” module, this implementation of module_a is enough.

Let's now take a look at the implemented unit test:

# tests/test_module_b.py

import unittest
from unittest.mock import patch, call

from custom_module.module_b import function_b

class TestModuleB(unittest.TestCase):
    def setUp(self):
        self.patch_builtins_print = patch("builtins.print")
        self.mock_builtins_print = self.patch_builtins_print.start()

    def tearDown(self):
        self.patch_builtins_print.stop()

    def test_should_print_function_b_msg(self):
        expected_calls = [
            call("I'm the function B")
        ]

        function_b()

        self.mock_builtins_print.assert_has_calls(expected_calls)

At the beginning of the file, we are importing, in the sequence:

unittest: a module dedicated to implementing the Python unit tests;
patch e call: functions that are used for mocking objects, functions and calls;
function_b: a function from module_b, which we want to test.

In the configuration function of our test scenarios, called setUp, we are mocking the native print python function (print) and the function_a of the module_a that is “broken”. This configuration function is executed before each implemented test scenario of the TestModuleB class.

The tearDown function is the function that is executed after each implemented test scenario of the TestModuleB class and it is responsible for resetting the mocks configured in the setUp function.

Our test scenario is the test_should_print_function_b_msg function and we define in it within the expected_calls variable which are the calls we expect from the print function and which parameter is passed to it.

The way our mini project is implemented, if we run our test scenario, we will have the following result:

$ python3 -m unittest discover -v
tests.test_module_b (unittest.loader._FailedTest) ... ERROR

======================================================================
ERROR: tests.test_module_b (unittest.loader._FailedTest)
----------------------------------------------------------------------
ImportError: Failed to import test module: tests.test_module_b
Traceback (most recent call last):
  File "/usr/lib/python3.10/unittest/loader.py", line 436, in _find_test_path
    module = self._get_module_from_name(name)
  File "/usr/lib/python3.10/unittest/loader.py", line 377, in _get_module_from_name
    __import__(name)
  File "/home/leogiraldimg/repos/mock-imported-python-modules-implementation/tests/test_module_b.py", line 4, in <module>
    from custom_module.module_b import function_b
  File "/home/leogiraldimg/repos/mock-imported-python-modules-implementation/custom_module/module_b.py", line 1, in <module>
    from custom_module.module_a import function_a
  File "/home/leogiraldimg/repos/mock-imported-python-modules-implementation/custom_module/module_a.py", line 4, in <module>
    raise Exception("some generic error in module a")
Exception: some generic error in module a

----------------------------------------------------------------------
Ran 1 test in 0.000s

FAILED (errors=1)

As we can see, the module_a exception is thrown and prevents our unit test from running successfully.

To fix this problem, we need to mock the module_a as follows:

import sys
import unittest
from unittest.mock import MagicMock, patch, call

sys.modules["custom_module.module_a"] = MagicMock()
from custom_module.module_b import function_b

class TestModuleB(unittest.TestCase):
    def setUp(self):
        self.patch_builtins_print = patch("builtins.print")
        self.mock_builtins_print = self.patch_builtins_print.start()

    def tearDown(self):
        self.patch_builtins_print.stop()

    def test_should_print_function_b_msg(self):
        expected_calls = [
            call("I'm the function B")
        ]

        function_b()

        self.mock_builtins_print.assert_has_calls(expected_calls)

The first step is to import the Python module sys and through it mock the import of the module_a file in sys.modules["custom_module.module_a"] = MagicMock(). Here we are using the MagicMock class from the mock unit testing utility.

Now, we have the following result:

$ python3 -m unittest discover -v
test_should_print_function_b_msg (tests.test_module_b.TestModuleB) ... ok

----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

We can further improve our code by moving the mock from module_a to a file dedicated only to these needs:

# tests/test_module_b.py

import unittest
from unittest.mock import MagicMock, patch, call

import tests.mock_modules
from custom_module.module_b import function_b

class TestModuleB(unittest.TestCase):
    def setUp(self):
        self.patch_builtins_print = patch("builtins.print")
        self.mock_builtins_print = self.patch_builtins_print.start()

    def tearDown(self):
        self.patch_builtins_print.stop()

    def test_should_print_function_b_msg(self):
        expected_calls = [
            call("I'm the function B")
        ]

        function_b()

        self.mock_builtins_print.assert_has_calls(expected_calls)

# tests/mock_modules.py

import sys
from unittest.mock import MagicMock

sys.modules["custom_module.module_a"] = MagicMock()