Pytest in the Networking World – Part 3

This blog post will be the final entry in my “Pytest in the Networking World” blog series. Part 1 can be found here and Part 2 can be found here. I recommend reading through both so you have the basic knowledge of how pytest works and some insight into parameterization. In this post, I will be going over various aspects of mocking.

Why Do We Need Mocking?

Keeping in the same vein as our last blog, let’s say we wrote a function using netmiko that queries our Cisco 3650 for its MAC address table so that we can use that information and pass it to our normalize_mac_address function. Sometimes, this call completes in under a second. Other times, it may take up to a few seconds. When we write tests for our code, do we want to possibly wait a few seconds for our call to the 3650 to complete? We definitely don’t. Unit tests should be fast so that developers are more likely to utilize them. Is there a way we could possibly fake the API call? There is, and it’s called mocking! When we mock a function, we make a “dummy” copy of that function where we can control the logic in the function and what that function returns. Let’s walk through an example of how we’d take advantage of mocking.

Updating Our Files

At the end of the last blog, our Python file contained the following:

>>> NON_NORMALIZED_MACS = ["aabb.ccdd.eeff", "0011.2233.4455", "aa11.bb22.cc33"]
>>> NORMALIZED_MACS = ["aa:bb:cc:dd:ee:ff", "00:11:22:33:44:55", "aa:11:bb:22:cc:33"]
>>>
>>> @pytest.mark.parametrize("param_non_normalized_mac, param_normalized_mac", list(zip(NON_NORMALIZED_MACS, NORMALIZED_MACS)))
>>> def test_normalize_mac_address_lists(param_non_normalized_mac, param_normalized_mac):
>>>     assert normalize_mac_address(param_non_normalized_mac) == param_normalized_mac

Using parameterization, we were able to test our normalize_mac_address function with multiple MAC addresses while ensuring each set of data was treated as its own test. Now, let’s add our function that will connect to a Cisco 3650 and get its MAC address table to the same file that the normalize_mac_address function is in. We’ll also add a function that will get the MAC address table and normalize it all at once. Our file should now look like this:

>>> from netmiko import ConnectHandler
>>> from datetime import datetime
>>>
>>> def normalize_mac_address(mac):
>>>     if mac.count(".") == 2:
>>>         mac = f"{mac[0:2]}:{mac[2:4]}:{mac[5:7]}:{mac[7:9]}:{mac[10:12]}:{mac[12:14]}"
>>>     return mac
>>>
>>> def get_mac_address_table():
>>>     start_time = datetime.now()
>>>
>>>     mydevice = {
>>>         "device_type": "cisco_ios",
>>>         "host": "YOUR DEVICE IP",
>>>         "username": "YOUR DEVICE USERNAME",
>>>         "password": "YOUR PASSWORD FOR THE CONNECTING USER"
>>>     }
>>>     command = "show mac address-table"
>>>
>>>     net_connect = ConnectHandler(**mydevice)
>>>     output = net_connect.send_command(command, use_textfsm=True)
>>>     net_connect.disconnect()
>>>
>>>     print(f"\nTime to run: {datetime.now() - start_time}")
>>>     return output
>>>
>>> print(get_mac_address_table())

For information on Netmiko and TextFSM, check here.

I included some code so we could time how long it takes to run this function. Let’s run it now.

[{'destination_address': '0100.0ccc.cccc', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0100.0ccc.cccd', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0000', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0001', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0002', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0003', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0004', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0005', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0006', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0007', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0008', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0009', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.000a', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.000b', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.000c', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.000d', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.000e', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.000f', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0010', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0180.c200.0021', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': 'ffff.ffff.ffff', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, {'destination_address': '0014.1c57.a488', 'type': 'DYNAMIC', 'vlan': '300', 'destination_port': 'Gi1/0/1'}, {'destination_address': '00c8.8bca.55c7', 'type': 'STATIC', 'vlan': '300', 'destination_port': 'Vl300'}, {'destination_address': '0014.1c57.a488', 'type': 'DYNAMIC', 'vlan': '400', 'destination_port': 'Gi1/0/1'}, {'destination_address': '00c8.8bca.55d2', 'type': 'STATIC', 'vlan': '400', 'destination_port': 'Vl400'}, {'destination_address': '0014.1c57.a488', 'type': 'DYNAMIC', 'vlan': '450', 'destination_port': 'Gi1/0/1'}, {'destination_address': '00c8.8bca.55fd', 'type': 'STATIC', 'vlan': '450', 'destination_port': 'Vl450'}, {'destination_address': '0014.1c57.a488', 'type': 'DYNAMIC', 'vlan': '500', 'destination_port': 'Gi1/0/1'}, {'destination_address': '00c8.8bca.55d0', 'type': 'STATIC', 'vlan': '500', 'destination_port': 'Vl500'}]

Time to run: 0:00:04.369203

This gets us a list of dictionaries each containing the destination address, type, vlan, and destination port for each entry in the table. On the last line, you can see it took 4.4 seconds to connect to the device, run the command, parse the output, and present it back to us. Next we’ll need to update our normalize_mac_address function to account for the new data structure. Let’s also add a function to combine getting the MAC address table and normalizing it all in one function.

>>> from netmiko import ConnectHandler
>>> from datetime import datetime
>>>
>>> def normalize_mac_address(macs):
>>>     for entry in macs:
>>>         if entry["destination_address"].count(".") == 2:
>>>             new_mac = f"{entry['destination_address'][0:2]}:{entry['destination_address'][2:4]}:{entry['destination_address'][5:7]}:{entry['destination_address'][7:9]}:{entry['destination_address'][10:12]}:{entry['destination_address'][12:14]}"
>>>             entry['destination_address'] = new_mac
>>>     return macs
>>>
>>> def get_mac_address_table():
>>>     start_time = datetime.now()
>>>
>>>     mydevice = {
>>>         "device_type": "cisco_ios",
>>>         "host": "YOUR DEVICE IP",
>>>         "username": "YOUR DEVICE USERNAME",
>>>         "password": "YOUR PASSWORD FOR THE CONNECTING USER"
>>>     }
>>>     command = "show mac address-table"
>>>
>>>     net_connect = ConnectHandler(**mydevice)
>>>     output = net_connect.send_command(command, use_textfsm=True)
>>>     net_connect.disconnect()
>>>
>>>     print(f"\nTime to run: {datetime.now() - start_time}\n")
>>>     return output
>>>
>>>
>>> def get_mac_table_and_normalize():
>>>     macs = get_mac_address_table()
>>>     normalized_macs = normalize_mac_address(macs)
>>>     return normalized_macs
>>>
>>> print(get_mac_table_and_normalize())

If you run this now, you’ll see that the MAC addresses in the dictionary have been normalized.

[{'destination_address': '0100.0ccc.cccc', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, ...,
 {'destination_address': '00c8.8bca.55d0', 'type': 'STATIC', 'vlan': '500', 'destination_port': 'Vl500'}]

Time to run: 0:00:04.849528

[{'destination_address': '01:00:0c:cc:cc:cc', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}, ...,
{'destination_address': '00:c8:8b:ca:55:d0', 'type': 'STATIC', 'vlan': '500', 'destination_port': 'Vl500'}]

Middle entries have been taken out for brevity.

Great! Our new function works as expected. Now we’ll turn toward updating our tests file. Currently it looks like:

NON_NORMALIZED_MACS = ["aabb.ccdd.eeff", "0011.2233.4455", "aa11.bb22.cc33"]
NORMALIZED_MACS = ["aa:bb:cc:dd:ee:ff", "00:11:22:33:44:55", "aa:11:bb:22:cc:33"]

@pytest.mark.parametrize("param_non_normalized_mac, param_normalized_mac", list(zip(NON_NORMALIZED_MACS, NORMALIZED_MACS)), ids=[x for x in NON_NORMALIZED_MACS])
def test_normalize_mac_address_lists(param_non_normalized_mac, param_normalized_mac):
    assert normalize_mac_address(param_non_normalized_mac) == param_normalized_mac

This test will currently fail, as we have changed the normalize_mac_address function. For the sake of staying on topic, we’ll comment out that test and focus on testing out get_mac_table_and_normalize. Here is our file with a basic test written for our new function and our old test commented out.

>>> import pytest
>>> from mac_address import normalize_mac_address, get_mac_table_and_normalize
>>>
>>> #NON_NORMALIZED_MACS = ["aabb.ccdd.eeff", "0011.2233.4455", "aa11.bb22.cc33"]
>>> #NORMALIZED_MACS = ["aa:bb:cc:dd:ee:ff", "00:11:22:33:44:55", "aa:11:bb:22:cc:33"]
>>> #
>>> #@pytest.mark.parametrize("param_non_normalized_mac, param_normalized_mac", list(zip(NON_NORMALIZED_MACS, NORMALIZED_MACS)), ids=[x for x in NON_NORMALIZED_MACS])
>>> #def test_normalize_mac_address_lists(param_non_normalized_mac, param_normalized_mac):
>>> #    assert normalize_mac_address(param_non_normalized_mac) == param_normalized_mac
>>>
>>> def test_get_mac_table_and_normalize():
>>>     macs = get_mac_table_and_normalize()
>>>     assert True

Let’s run our test and look at the output.

============================================================= test session starts ==============================================================================
platform linux -- Python 3.8.5, pytest-6.2.4, py-1.10.0, pluggy-0.13.1
rootdir: /home/adam/repo/Sandbox/blog_pytest
collected 1 item

test_mac_address.py .                                                                                                                                      [100%]

============================================================== 1 passed in 8.95s ================================================================================

The test passes (because we are currently just asserting True), but you’ll notice that on the last line it states that our test passed in 8.95 seconds. We know that connecting to the device and retrieving the MAC address table is taking roughly 4 seconds, so let’s jump into how we’d mock this function.

Mocking

Let me update our function to leverage mocking, and then I’ll go over it.

>>> import pytest
>>> import re
>>> from mac_address import normalize_mac_address, get_mac_table_and_normalize
>>> from unittest.mock import patch
>>>
>>> @patch('mac_address.get_mac_address_table')
>>> def test_get_mac_table_and_normalize(mock_get_mac_address_table):
>>>     mock_get_mac_address_table.return_value = [{'destination_address': '01:00:0c:cc:cc:cc', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}]
>>>     macs = get_mac_table_and_normalize()
>>>     normalized_mac = macs[0].pop("destination_address")
>>>
>>>     assert re.match(r'([0-9a-fA-F]{4}[.]){2}([0-9a-fA-F]{4})', normalized_mac)
>>>     mock_get_mac_address_table.assert_called_once()

Commented test left out for brevity’s sake.

The first thing you may notice is that there are two new imports, import re and from unittest.mock import patch. The re library is not a necessary import to use mocks. I imported it for use in an assert statement that I’ll go over shortly. The unittest.mock import patch is necessary. This gives us access to the @patch() decorator that you see implemented on line 6. This decorator is what is actually doing the mocking. To use the decorator, you pass in the object that you want to mock using the syntax @patch(module.functionA) or @patch(module.ClassA). In this case, when I’m testing get_mac_table_and_normalize, I want to mock the get_mac_address_table function that gets called in my get_mac_table_and_normalize function. I implement my decorator using @patch('mac_address.get_mac_address_table'); and in my test function definition, I add the argument mock_get_mac_address_table. When you mock an object, that mocked object gets passed into the decorated function as an argument so you need to account for that. I typically name the function with mock_ prepended to it.

On line 8, I define what I want the return value to be when my mocked function is called. I can make this whatever I want, but you typically would make this similar to what you would expect from the original function. Line 9, I call our get_mac_table_and_normalize function. Line 10, I get the value for the destination_address key from the first dictionary in the returned list from calling get_mac_table_and_normalize and store it in a variable. In lines 12 and 13, I have two assertions taking place. The first ensures that the value I popped from our returned list conforms to our normalized MAC address format by leveraging regex. The second asserts that our mock function was called only one time. Let’s run our newly written test and check out the results.

=========================================================== test session starts =======================================================================================
platform linux -- Python 3.8.5, pytest-6.2.4, py-1.10.0, pluggy-0.13.1 -- /home/adam/.virtualenvs/blog/bin/python3
cachedir: .pytest_cache
rootdir: /home/adam/repo/Sandbox/blog_pytest
collected 1 item

test_mac_address.py::test_get_mac_table_and_normalize PASSED                                                                                                      [100%]

========================================================== 1 passed in 4.59s ===========================================================================================

You’ll notice our test passes and that it now took 4.59 seconds to complete. That is exactly what we expected as we mocked our call to our network device which took roughly 4 seconds. To provide another example, let’s say we wanted to mock the normalize_mac_address function. How would we go about that? We can follow the same process we did for the get_mac_address_table function.

>>> import pytest
>>> import re
>>> from mac_address import normalize_mac_address, get_mac_table_and_normalize
>>> from unittest.mock import patch
>>>
>>> @patch('mac_address.get_mac_address_table')
>>> @patch('mac_address.normalize_mac_address')
>>> def test_get_mac_table_and_normalize(mock_normalize_mac_address, mock_get_mac_address_table):
>>>     mock_get_mac_address_table.return_value = [{'destination_address': '01:00:0c:cc:cc:cc', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}]
>>>     mock_normalize_mac_address.return_value = [{'destination_address': '0100.0ccc.cccc', 'type': 'STATIC', 'vlan': 'All', 'destination_port': 'CPU'}]
>>>     macs = get_mac_table_and_normalize()
>>>     normalized_mac = macs[0].pop("destination_address")
>>>
>>>     assert re.match(r'([0-9a-fA-F]{4}[.]){2}([0-9a-fA-F]{4})', normalized_mac)
>>>     mock_get_mac_address_table.assert_called_once()

We are now mocking both functions that are called when we call get_mac_table_and_normalize. We added a patch decorator and included mac_address.normalize_mac_address in its parameter. We also have to add another argument to our test function definition. When mocking more than one object, the lowest patch decorator is mapped to the leftmost argument in the function definition. If we run this, we get the following output:

======================================================== test session starts =============================================================================================
platform linux -- Python 3.8.5, pytest-6.2.4, py-1.10.0, pluggy-0.13.1
rootdir: /home/adam/repo/Sandbox/blog_pytest
collected 1 item

test_mac_address.py .                                                                                                                                                [100%]

======================================================== 1 passed in 4.45s ================================================================================================

Our test again passes. This time we have only a slight decrease in test time to 4.45 seconds. This was expected as the normalize_mac_address function is not doing anything extensive, so mocking it didn’t yield much decrease in computation time.

To prove the test is using the mocked return values, you can change the line 10 destination address to something that doesn’t conform to our normalized MAC format and the test will fail.

Mocking Tips

  • There are different implementations of mocking in pytest. I chose the above method because I felt it was the easiest to see and wrap your head around. To see the other implementations of mocking, you can check out the documentation here.
  • When referring to the object you want to mock in the decorator argument, you need to think about mocking the object where it is looked up rather than where it may be defined. This documentation goes more in-depth.
  • You can do a wide variety of things with a mocked object. You can manipulate it to return whatever you want (as shown in our example) or have it return dynamic results using side effects. More information on that and other ways you can manipulate mock functions can be found here.
  • A multitude of assertions can be made on mocked functions. The documentation here goes over them.

Conclusion

This blog is just the tip of the iceberg when it comes to mocking. It’s a pretty extensive feature and can get very complicated very quickly. While this application of mocking may not be the most practical, hopefully these examples were easy enough to follow along with and you were able to gather what mocking is and one of its applications. If you are interested in looking into more robust applications, you can check out the tests for our pyntc repository here. If you have any questions or want to discuss any tests we at NTC may have written that you’d like clarification on, definitely come by and reach out to us on our NTC community Slack.

-Adam



ntc img
ntc img

Contact Us to Learn More

Share details about yourself & someone from our team will reach out to you ASAP!

Author