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-levelnetboxobject is already replaced withnetbox_mockand the return value of the call is not really used so it worked also without the specific mock) side_effectis 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 thedescription="Prod"filter), and the second item (with subnets 2 and 3) is used with the second call (when using therole_id=role.idfilter). 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!