Mocking imported Python modules implementation
Tags: python, unittest
Published: 2022-09-19 00:32:05

---
title: Mocking imported Python modules implementation
published: true
description: We're are going to learn how to "physically" mock a module file that we want to ignore in our unit tests implementation.
tags: python, unittest
# cover_image: https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcTNkVF6WI-8Vn1iurUowQDHabMad43PIilrUg&usqp=CAU
# Use a ratio of 100:42 for best results.
---
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`](https://github.com/ansible/ansible/blob/devel/lib/ansible/module_utils/basic.py):

```python
# 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:

```python
├── 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:

```python
# 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:

```python
# 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:

```python
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:

```python
# 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()
```