Python unit testing: mocking with patch decorator

A piece of code rarely does all the work itself. It calls another code which calls something else.

Why that's not good for unit testing:

• Test results could change on each test run, for example if email isn't sent due network error
• Complex setup is required: database setup and test data for testing SQL query
• Long call chains take a lot of time

So for testing would be better to replace another (let's call them external) function calls with calls that do nothing or return predefined results: for example, do nothing on send email or return predefined data on the database SQL query. It's called mocking external calls.

1. Python Unit test: mocking for test method

Standard Python mocking library allows replacing external calls with mock functions in the test code without changing the code under testing:

Code to test, reports.py:

def send_mail(msg):
  print(f'send mail: {msg}')
  

def send_report(rows):
  result = sum(rows)
  send_mail(f'Report sum: {result}')
  return result

Unit Test, we'll apply patch decorator to the test method to mock the send_mail function in the reports module:

from unittest import TestCase, main
from unittest.mock import patch

from reports import send_report


class TestReport(TestCase):

  @patch('reports.send_mail')
  def test_send_report(self, send_mail_mock):
    rows = [1, 2, 3]
    result = send_report(rows)
    self.assertEqual(sum(rows), result)  # check the result
    send_mail_mock.assert_called_once()  # check the send_mail call is made
    
    expected_str = f'Report sum: {result}'
    send_mail_mock.assert_called_once_with(expected_str)  # check call parameter
    
    
if __name__ == '__main__':
  main()

The test checks that the function result and the send_mail is called once with the correct parameters.

The patch decorator for test function does:

• Creates mock function
• Changes the send_mail reference in the reports module to this function: the patch argument is path to the patched function, module and its name
• Passes the mock function into the test function as send_mock parameter

2. Mocking for code block in test method

How to patch without changing the test method declaration. Use with patch code block as context manager:

class TestReport(TestCase):

  def test_send_report(self):
    rows = [1, 2, 3]
    
    with patch('reports.send_mail') as send_mail_mock:
      result = send_report(rows)
      self.assertEqual(sum(rows), result)
      send_mail_mock.assert_called_once()

It works the same as the option 1.

3. Mock function return value

Let's consider the case the external call is made to calculate the return value:

Code to test, reports.py:

def calculate_sum():
  rows = query_database()  # returns integer array
  return sum(rows)

Mock return_value property is set in Unit test to the value that should be returned from the mock function:

class TestReport(TestCase):

  @patch('reports.query_database')
  def test_calculate_sum(self, query_db_mock):
    rows = [1, 2, 3]
    
    query_db_mock.return_value = rows  # return value for mock function is set
    result = calculate_sum()
    self.assertEqual(sum(rows), result)
    query_db_mock.assert_called_once()

4. Error testing

How to test that errors from external calls are handled correctly. Let's raise ConnectionError in the mock function setting the exception class to the side_effect property:

class TestReport(TestCase):

  # mock query_database function imported in reports module
  @patch('reports.query_database')
  def test_connection_error(self, query_db_mock):
    query_db_mock.side_effect = ConnectionError
    with self.assertRaises(ConnectionError)
        result = calculate_sum()
        query_db_mock.assert_called_once()

5. Mock for all test methods in class

What if there are many tests which require patching the same function? To avoid patch on every test method, apply patch decorator on the test class instead and add the mock argument to every test method:

# mock query_database function for every test method
@patch('reports.query_database')
class TestReport(TestCase):

  def test_calculate_sum(self, query_db_mock):
    rows = [1, 2, 3]
    
    query_db_mock.return_value = rows
    result = calculate_sum()
    self.assertEquals(sum(rows), result)
    query_mock.assert_called_once()
    
  def test_connection_error(self, query_db_mock):
    query_db_mock.side_effect = ConnectionError
    with self.assertRaises(ConnectionError)
        result = calculate_sum()
        query_db_mock.assert_called_once()

6. Patch call chain, mock return values: MagicMock

How to return mock object from mocked function and patch mock object methods:

reports.py:

def get_report():
   print("get report object")
   

def calculate_sum():
  report = get_report()
  rows = report.get_values()
  return sum(rows)

MagicMock is used for return value in the Unit test, the MagicMock properties and functions are mocks as well:

from unittest.mock import MagicMock

 
class TestReport(TestCase):
  
  @patch('reports.get_report')
  def test_calculate_sum(self, get_report_mock):
    rows = [1, 2, 3]
    
    report_mock = MagicMock()
    report_mock.get_values.return_value = rows  # patch call chain
    
    get_report_mock.return_value = report_mock  # return value is mock
    
    result = calculate_sum()
    self.assertEqual(sum(rows), result)
    report_mock.get_values.assert_called_once()  # check mock method is called once
    get_report_mock.assert_called_once()

Please note, we created a test for get_report() function even without implementing it yet.