Skip to main content

Lets unittest using Python Mock, wait, but what to Mock?

At Senic, on our Hub, for managing applications, we use Supervisord. I am not sure about its python3 compatibility, but it is one of the reason we still have dependency on Python2.7 Given that Python2.7 life support clock is ticking, we recently merged big changes to use Systemd instead. I came across this small, clean, Python API for managing systemd services. We included it in our stack and I wrote a small utility function for it:

import logging
from sysdmanager import SystemdManager
import dbus.exceptions


logger = logging.getLogger(__name__)

def manage_service(service_name, action):
    '''This function is to manage systemd units passed to it in
    service_name argument. It will try to stop/start/restart unit
    based on value passed in action.
    '''
    try:
	systemd_manager = SystemdManager()
    except dbus.exceptions.DBusException:
	logger.exception("Systemd service is not accessible via DBus")
    else:
	if action == "start":
	    if not systemd_manager.start_unit(service_name):
		logger.info("Failed to start {}".format(service_name))
	elif action == "stop":
	    if not systemd_manager.stop_unit(service_name):
		logger.info("Failed to stop {}".format(service_name))
	elif action == "restart":
	    if not systemd_manager.restart_unit(service_name):
		logger.info("Failed to restart {}".format(service_name))
	else:
	    logger.info(
		"Invalid action: {} on service: {}".format(action, service_name)
	    )

With this in place, manage_service can be imported in any module and restarting any service is , manage_service('service_name', 'restart'). Next thing was putting together some unittests for this utility, to confirm if its behaving the way it should.

This smallish task got me confused for quite some time. My first doubt was how and where to start? Library needed SystemBus DBus, Systemd's DBus API to start, stop, load, unload systemd services. I can't directly write tests against these APIs as they would need root privilege to work, additionally, they won't work on travis. So, I realized, I will need to mock, and with that realization came second doubt, which part to mock? Should I mock things needed by library or should I mock library? When I looked for mocking Systemd APIs on DBus via dbus-mock, I realized this can become too big of task. So lets mock library object and functions which gets called when I call the utility function manage_service. I had read/noticed python's mock support, and while trying to understand it, it came across as a powerful tool and I remembered Uncle Ben has once rightly said, with great power comes great responsibility. At one point, I was almost convinced of hijacking the utility function and having asserts around different branching happening there. But soon I also realized it will defeat the purpose of unit-testing the utility and sanity prevailed. After looking around at lots of blogs, tutorials and conversations with peers, I carefully mocked some functions from SystemdManager, like stop_unit, start_unit, which gets internally called from the library and that way I was able to write tests for different arguments which could be passed to manage_service. At the end the tests looked something like this:

import unittest
from systemd_process_utils import manage_service
from systemd_process_utils import SystemdManager
from unittest import mock

class TestSystemdUtil(unittest.TestCase):
    service_name = "service_name"
    @mock.patch('senic_hub.commons.systemd_process_utils.SystemdManager')
    def test_manage_service(self, mock_systemd):
	# When: start_unit works, it returns True
	mock_systemd.return_value.start_unit.return_value = True
	manage_service(self.service_name, "start")
	mock_systemd().start_unit.assert_called_with(self.service_name)

	# When: start_unit fails, returns False
	mock_systemd.return_value.start_unit.return_value = False
	manage_service(self.service_name, "start")
	mock_systemd().start_unit.assert_called_with(self.service_name)

	# When: stop_unit works, returns True
	mock_systemd.return_value.stop_unit.return_value = True
	manage_service(self.service_name, "stop")
	mock_systemd().stop_unit.assert_called_with(self.service_name)

if __name__ == '__main__':
    unittest.main()