More Mocking with Pytest

Programming
Comments

In my Basic Mocking with Pytest post I showed simple cases. Now let’s do more mocking.

First I’ll create the environment for my code:

$ mkdir more-mocking
$ cd more-mocking
$ python3 -m venv venv
$ . venv/bin/activate
(venv)$ pip install -U pip wheel
...
Successfully installed pip-24.0 wheel-0.43.0
(venv)$ pip install pynetbox pytest
...
Successfully installed certifi-2024.2.2 charset-normalizer-3.3.2 idna-3.7 iniconfig-2.0.0 packaging-23.2 pluggy-1.5.0 pynetbox-7.3.3 pytest-8.2.0 requests-2.31.0 urllib3-2.2.1
(venv)$

My application files:

# site_utils.py

import pynetbox


def get_netbox(url: str, token: str) -> pynetbox.api:
    return pynetbox.api(url, token)


def get_prod_subnets(netbox: pynetbox.api, sitename: str) -> list[str]:
    site = netbox.dcim.sites.get(name=sitename)
    if not site:
        raise ValueError(f"Site {sitename} not found")
    # Get Prod prefixes by description
    prefixes = netbox.ipam.prefixes.filter(
        site_id=site.id,
        description="Prod",
    )
    prod_subnets = [p.prefix for p in prefixes]
    return prod_subnets
# my_app.py

import site_utils


# In reality the URL and token come from a config file,
# environment variable, etc
netbox = site_utils.get_netbox(
    "http://netbox-test.lein.io/",
    token="176d4c04ccc8f2a549ea6fd393567d9da5a796ff",
)
sitename = "TestSite"

subnets = site_utils.get_prod_subnets(netbox, sitename)

print(f"Prod subnets in {sitename}: {', '.join(subnets)}")

Let’s run this:

(venv)$ python my_app.py
Prod subnets in TestSite: 10.10.1.0/24, 10.10.2.0/24
(venv)$

Nice, it works: it retrieves the prefixes from my NetBox server.

Tests, round 1

Let’s write some tests for the functions in site_utils module:

# test_site_utils.py

from unittest.mock import Mock

import pynetbox
import pytest

import site_utils


def test_get_netbox():
    netbox = site_utils.get_netbox(
        "https://netbox.example.com/",
        token="1234567890",
    )
    assert isinstance(netbox, pynetbox.api)


def test_get_prod_subnets():
    # Setup the mocks
    netbox_mock = Mock()
    netbox_mock.dcim.sites.get.return_value = Mock(id=123)
    netbox_mock.ipam.prefixes.filter.return_value = [
        Mock(prefix="10.20.31.0/24"),
        Mock(prefix="10.20.32.0/24"),
    ]
    # Now call the tested function
    subnets = site_utils.get_prod_subnets(
        netbox_mock,
        "My Site",
    )
    # Check the results
    assert isinstance(subnets, list)
    assert subnets == ["10.20.31.0/24", "10.20.32.0/24"]


def test_get_prod_subnets_invalid_site():
    netbox_mock = Mock()
    netbox_mock.dcim.sites.get.return_value = None
    with pytest.raises(ValueError):
        subnets = site_utils.get_prod_subnets(
            netbox_mock,
            "Non-existent Site",
        )

Let’s look at this test code. First, it imports the Mock class from Python standard library, the pynetbox and pytest modules, and our own site_utils module.

Note: If you are into using pytest-mock plugin, you can skip importing Mock, and then use the mocker fixture instead:

def test_something(mocker):
    my_mock = mocker.Mock()
    ...

The first test, test_get_netbox(), is straightforward: it just instantiates a pynetbox.api object and checks the instance type. pynetbox does not make any external calls in that phase so there is nothing to mock yet.

The second test, test_get_prod_subnets(), does some mocking. It creates a mock that simulates the netbox.dcim.sites.get() and netbox.ipam.prefixes.filter() calls, the return values of the calls to be exact. Mocking is quite easy in this case because the NetBox instance is specified as the function argument. This is called dependency injection: the function gets the necessary execution information as arguments instead of depending on some fixed internal logic.

The third test, test_get_prod_subnets_invalid_site(), is very similar, but it simulates the case where the netbox.dcim.sites.get() call returns None and thus the function raises ValueError.

Let’s run the tests:

(venv)$ pytest -v
======================== test session starts =========================
platform linux -- Python 3.11.2, pytest-8.2.0, pluggy-1.5.0 -- /home/markku/more-mocking/venv/bin/python3
cachedir: .pytest_cache
rootdir: /home/markku/more-mocking
collected 3 items

test_site_utils.py::test_get_netbox PASSED                     [ 33%]
test_site_utils.py::test_get_prod_subnets PASSED               [ 66%]
test_site_utils.py::test_get_prod_subnets_invalid_site PASSED  [100%]

========================= 3 passed in 0.10s ==========================
(venv)$

Success, the mocking works and no actual NetBox connections were used in the tests.

Tests, round 2

Now, someone notices the nice feature of prefix/VLAN roles in NetBox, and starts using them in new prefixes. The old prefixes still use the description field. Our get_prod_subnets() code definitely needs modifications for this. Let’s change it:

# site_utils.py

import pynetbox


def get_netbox(url: str, token: str) -> pynetbox.api:
    return pynetbox.api(url, token)


def get_prod_subnets(netbox: pynetbox.api, sitename: str) -> list[str]:
    site = netbox.dcim.sites.get(name=sitename)
    if not site:
        raise ValueError(f"Site {sitename} not found")
    # Get Prod prefixes by description
    prefixes = netbox.ipam.prefixes.filter(
        site_id=site.id,
        description="Prod",
    )
    prod_subnets = [p.prefix for p in prefixes]
    # Get Prod prefixes by role
    role = netbox.ipam.roles.get(name="Prod")
    prefixes = netbox.ipam.prefixes.filter(
        site_id=site.id,
        role_id=role.id,
    )
    prod_subnets += [p.prefix for p in prefixes]
    return prod_subnets

The difference is that now the code also gets the Prod role and then filters the prefixes with that, adds the resulted subnets to the prod_subnets list, and returns the list.

Let’s run the app code again:

(venv)$ python my_app.py
Prod subnets in TestSite: 10.10.1.0/24, 10.10.2.0/24, 10.10.3.0/24
(venv)$

Now there is also the third subnet shown as it has the Prod role set.

How about our tests? Let’s try running them again:

(venv)$ pytest -v
======================== test session starts =========================
platform linux -- Python 3.11.2, pytest-8.2.0, pluggy-1.5.0 -- /home/markku/more-mocking/venv/bin/python3
cachedir: .pytest_cache
rootdir: /home/markku/more-mocking
collected 3 items

test_site_utils.py::test_get_netbox PASSED                     [ 33%]
test_site_utils.py::test_get_prod_subnets FAILED               [ 66%]
test_site_utils.py::test_get_prod_subnets_invalid_site PASSED  [100%]

============================== FAILURES ==============================
_______________________ test_get_prod_subnets ________________________

    def test_get_prod_subnets():
        # Setup the mocks
        netbox_mock = Mock()
        netbox_mock.dcim.sites.get.return_value = Mock(id=123)
        netbox_mock.ipam.prefixes.filter.return_value = [
            Mock(prefix="10.20.31.0/24"),
            Mock(prefix="10.20.32.0/24"),
        ]
        # Now call the tested function
        subnets = site_utils.get_prod_subnets(
            netbox_mock,
            "My Site",
        )
        # Check the results
        assert isinstance(subnets, list)
>       assert subnets == ["10.20.31.0/24", "10.20.32.0/24"]
E       AssertionError: assert ['10.20.31.0/...0.20.32.0/24'] == ['10.20.31.0/...0.20.32.0/24']
E
E         Left contains 2 more items, first extra item: '10.20.31.0/24'
E
E         Full diff:
E           [
E               '10.20.31.0/24',
E               '10.20.32.0/24',...
E
E         ...Full output truncated (3 lines hidden), use '-vv' to show

test_site_utils.py:32: AssertionError
====================== short test summary info =======================
FAILED test_site_utils.py::test_get_prod_subnets - AssertionError: assert ['10.20.31.0/...0.20.32.0/24'] == ['10.20.31.0/......
==================== 1 failed, 2 passed in 0.12s =====================
(venv)$

Oh no, the test is failing now. Why is that?

It is because of the mocking of the netbox.ipam.prefixes.filter() call results: In the new code the call is made twice, and thus the mock is used twice, causing the same subnets to be duplicated in the results.

netbox_mock.ipam.prefixes.filter.return_value is a static list, so some modifications are needed to be able to get different results for the different calls in the code. Ideally, we should also mock the result of the netbox.ipam.roles.get() call.

Let’s use side_effect instead of return_value in the prefix filtering calls:

# test_site_utils.py

from unittest.mock import Mock

import pynetbox
import pytest

import site_utils


def test_get_netbox():
    netbox = site_utils.get_netbox(
        "https://netbox.example.com/",
        token="1234567890",
    )
    assert isinstance(netbox, pynetbox.api)


def test_get_prod_subnets():
    # Define our subnets to use:
    SUBNETS = [
        "10.20.31.0/24",
        "10.20.32.0/24",
        "10.20.33.0/24",
        "10.20.34.0/24",
    ]
    # Setup the mocks
    netbox_mock = Mock()
    netbox_mock.dcim.sites.get.return_value = Mock(id=123)
    netbox_mock.ipam.roles.get.return_value = Mock(id=456)
    netbox_mock.ipam.prefixes.filter.side_effect = [
        [Mock(prefix=SUBNETS[0]), Mock(prefix=SUBNETS[1])],
        [Mock(prefix=SUBNETS[2]), Mock(prefix=SUBNETS[3])],
    ]
    # Now call the tested function
    subnets = site_utils.get_prod_subnets(
        netbox_mock,
        "My Site",
    )
    # Check the results
    assert isinstance(subnets, list)
    assert subnets == SUBNETS


def test_get_prod_subnets_invalid_site():
    netbox_mock = Mock()
    netbox_mock.dcim.sites.get.return_value = None
    with pytest.raises(ValueError):
        subnets = site_utils.get_prod_subnets(
            netbox_mock,
            "Non-existent Site",
        )

Let’s first see what was changed:

  • There is a list of all the subnets that are used in the test, to avoid typing the same prefixes again
  • There is a specific mock for the netbox.ipam.roles.get() call, even though it doesn’t do much here yet (the upper-level netbox object is already replaced with netbox_mock and the return value of the call is not really used so it worked also without the specific mock)
  • side_effect is used when mocking the list of returned prefixes from NetBox: It is a list of lists, where the first item of the list (the one with subnets 0 and 1) is used when the mock is first called (when using the description="Prod" filter), and the second item (with subnets 2 and 3) is used with the second call (when using the role_id=role.id filter). Thus we get different return values for the different .filter() calls in the test.

But does it work?

(venv)$ pytest -v
======================== test session starts =========================
platform linux -- Python 3.11.2, pytest-8.2.0, pluggy-1.5.0 -- /home/markku/more-mocking/venv/bin/python3
cachedir: .pytest_cache
rootdir: /home/markku/more-mocking
collected 3 items

test_site_utils.py::test_get_netbox PASSED                     [ 33%]
test_site_utils.py::test_get_prod_subnets PASSED               [ 66%]
test_site_utils.py::test_get_prod_subnets_invalid_site PASSED  [100%]

========================= 3 passed in 0.10s ==========================
(venv)$

Yes, it works.

Tests, round 3

If we take even more deeper dive into the side_effect attribute, we can also do dynamic return values, based on the arguments used in the calls:

# test_site_utils.py (continued)

def test_get_prod_subnets_with_dynamic_mock():
    # Define our subnets to use:
    SUBNETS = [
        "10.20.31.0/24",
        "10.20.32.0/24",
        "10.20.33.0/24",
        "10.20.34.0/24",
    ]

    def _filter_call(**kwargs):
        if kwargs.get("description", "") == "Prod":
            return [
                Mock(prefix=SUBNETS[0]),
                Mock(prefix=SUBNETS[1]),
            ]
        elif kwargs.get("role_id", 0) == 456:
            return [
                Mock(prefix=SUBNETS[2]),
                Mock(prefix=SUBNETS[3]),
            ]

    # Setup the mocks
    netbox_mock = Mock()
    netbox_mock.dcim.sites.get.return_value = Mock(id=123)
    netbox_mock.ipam.roles.get.return_value = Mock(id=456)
    netbox_mock.ipam.prefixes.filter.side_effect = _filter_call
    # Now call the tested function
    subnets = site_utils.get_prod_subnets(
        netbox_mock,
        "My Site",
    )
    # Check the results
    assert isinstance(subnets, list)
    assert subnets == SUBNETS

We added an internal _filter_call() function that is called whenever the mocked function is called, with the arguments of the original calls. Thus we can check the call arguments and then decide the returned values dynamically.

To verify:

(venv)$ pytest -v
======================== test session starts =========================
platform linux -- Python 3.11.2, pytest-8.2.0, pluggy-1.5.0 -- /home/markku/more-mocking/venv/bin/python3
cachedir: .pytest_cache
rootdir: /home/markku/more-mocking
collected 4 items

test_site_utils.py::test_get_netbox PASSED                     [ 25%]
test_site_utils.py::test_get_prod_subnets PASSED               [ 50%]
test_site_utils.py::test_get_prod_subnets_invalid_site PASSED  [ 75%]
test_site_utils.py::test_get_prod_subnets_with_dynamic_mock PASSED [100%]

========================= 4 passed in 0.10s ==========================
(venv)$

Success!

Leave a Reply