Python Unittest: A Guide to Patching, Mocking, and MagicMocks

Python’s unittest framework is a powerful tool for writing and running unit tests for your code. In this article, we'll explore some advanced techniques in Python unit testing, specifically focusing on patch, mock, and MagicMock. These features allow you to isolate and control the behaviour of your code during testing, making your tests more robust and accurate.

Introduction: In software development, testing is a crucial part of ensuring the correctness and reliability of your code. Python provides the unittest module, that allows you to write unit tests to check the behaviour of your functions and classes. However, your code often interacts with external dependencies like databases, APIs, or third-party libraries. You need to isolate your code from these external factors to create meaningful unit tests in such scenarios. This is where mocking and patching come into play.

Setting Up the Environment: Before diving into unit testing, ensure you have a suitable environment set up. You should have Python installed, and you may want to create a virtual environment to manage dependencies. You’ll also need the unittest.mock module, which is included in Python 3.3 and later. You can install it using pip if needed.

pip install unittest-mock

The Basics of Unit Testing: Before we explore mocking, it’s essential to understand the basics of unit testing in Python. You create test cases by creating subclasses of unittest.TestCase and using various assertion methods such as assertEqual, assertTrue, and assertRaises to verify the behaviour of your code.

import unittest
def add(a, b):
    return a + b
class TestMathOperations(unittest.TestCase):
    def test_addition(self):
        result = add(2, 3)
        self.assertEqual(result, 5)
if __name__ == '__main__':
    unittest.main()

For a detailed overview of unit tests in Python please check my other blog: https://readmedium.com/mastering-test-automation-with-pythons-unittest-framework-a95d35613e25

For running the unit test you can use python -m unittest <unitest_file.py> or if you are using python3 then python3 -m unittest <unitest_file.py> .

The Need for Mocking: As mentioned earlier, real-world applications often rely on external resources and dependencies. During unit testing, you want to focus on the specific piece of code you’re testing, not the external systems or functions it interacts with. Mocking is the process of replacing these external dependencies with simulated objects that allow you to control their behaviour.

Let’s delve into the three key techniques: patch, mock, and MagicMock but before that, you can set the Python path as your project path so that you won’t get the module not found error.

export PYTHONPATH=/your_project_path/

Using patch for Function Mocking: patch is a powerful decorator that temporarily replaces functions or objects in your code with mock objects. It's especially useful when you need to mock a function that's called within the code you're testing.

import unittest
from unittest.mock import patch
def send_notification(email):
    # Real implementation to send an email
    pass
def process_data(data, email):
    # Process data and send a notification
    send_notification(email)
class TestProcessData(unittest.TestCase):
    @patch('your_module.send_notification')
    def test_process_data(self, mock_send_notification):
        # Replace send_notification with a mock object
        mock_send_notification.return_value = True
        # Call the function under test
        result = process_data("Some data", "[email protected]")
        # Assert that send_notification was called with the expected arguments
        mock_send_notification.assert_called_with("[email protected]")
if __name__ == '__main__':
    unittest.main()

When using @patch('your_module.send_notification'), if you set the python_path to your project directory, your_module should correspond to the file path where the send_notification function is defined. In the given example, if you define send_notification within the test file, then your_module would be the name of the test file.

Let’s check another example where you need to fetch the value stored in the Airflow variables.

import unittest
from airflow.models import Variable
from unittest.mock import patch
def airflow_variable_value(key):
    # Fetching Airflow Variable
    return Variable.get(key)
class TestAirflowVariable(unittest.TestCase):
    @patch('your_module.Variable')
    def test_process_data(self, mock_airflow_varibale):
        # Replace Variable.get with a mock object
        mock_airflow_varibale.get.return_value = 'Value from airflow variables'
        # Call the function under test
        result = airflow_variable_value('key')
        # Assert that airflow_variable_value returns the expected arguments
        self.assertEqual(result,'Value from airflow variables')
if __name__ == '__main__':
    unittest.main()

Using patch.object for Method Mocking: If you want to mock a method of an object, you can use patch.object. This allows you to replace the method within a class temporarily.

import unittest
from unittest.mock import patch
class MyService:
    def send_notification(self, email):
        # Real implementation to send an email
        pass

def process_data(service, data, email):
    # Process data and send a notification using the service object
    service.send_notification(email)

class TestProcessData(unittest.TestCase):
    @patch.object(MyService, 'send_notification')
    def test_process_data(self, mock_send_notification):
        # Replace send_notification method with a mock
        mock_send_notification.return_value = True
        # Create an instance of MyService
        service = MyService()
        # Call the function under test
        result = process_data(service, "Some data", "[email protected]")
        # Assert that send_notification was called with the expected arguments
        mock_send_notification.assert_called_with("[email protected]")
if __name__ == '__main__':
    unittest.main()

Here is an example of patch with the with constructor

import unittest
from unittest.mock import patch, MagicMock
from unittest.mock import Mock

def send_notification(email):
    # Real implementation to send an email
    pass

def process_data(data, email):
    # Process data and send a notification
    send_notification(email)

class TestProcessData(unittest.TestCase):
    def test_process_data(self):
        # Patch the function and return the actual value
        with patch('your_module.send_notification') as mock:
            mock.return_value = None
            # Call the function under test
            result=process_data('Some Data', '[email protected]')
            
            assert result == None

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

Mocking with unittest.mock:The unittest.mock module provides the Mock class for creating mock objects that replace functions and methods. This is a more manual way of creating mocks compared to using patch. You can control the behaviour of the mock and assert how it's used in your code.

import unittest
from unittest.mock import patch, MagicMock
from unittest.mock import Mock

class MyService:
    def send_notification(self, email):
        # Real implementation to send an email
        pass
def process_data(service, data, email):
    # Process data and send a notification using the service object
    service.send_notification(email)

class TestProcessData(unittest.TestCase):
    def test_process_data(self):
        # Create a mock object to replace send_notification method
        mock_send_notification = Mock()
        # Replace send_notification method with the mock object
        service = MyService()
        service.send_notification = mock_send_notification
        # Call the function under test
        result = process_data(service, "Some data", "[email protected]")
        # Assert that send_notification was called with the expected arguments
        mock_send_notification.assert_called_with("[email protected]")

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

MagicMocks: Magic methods in Python are special methods with double underscores at the beginning and end of their names (e.g., __init__, __str__, __getitem__). They allow you to customize the behaviour of objects or classes. MagicMocks are used for mocking classes and their magic methods in unit tests.

Understanding Magic Methods: Magic methods are an integral part of Python’s object-oriented programming. For example, the __str__ method is called when you use the str() function or print() to get a string representation of an object. You can customize these methods to define how your objects behave in certain situations.

class MyClass:
    def __str__(self):
        return "Original __str__ implementation"
obj = MyClass()
print(obj)  # Calls obj.__str__() and prints "Original __str__ implementation"

Using MagicMock to Mock Magic Methods: MagicMock is specifically designed to mock these magic methods. You can use it to replace the behaviour of magic methods in your classes during testing.

from unittest.mock import MagicMock
import unittest

class MyClass:
    def __str__(self):
        return "Original __str__ implementation"
def function_using_str_method(obj):
    return str(obj)

class TestFunctionUsingStrMethod(unittest.TestCase):

    def test_function_using_str_method(self):
        # Create a MagicMock object to replace the __str__ method of MyClass
        mock_obj = MagicMock(spec=MyClass)
        mock_obj.__str__.return_value = "Mocked __str__ implementation"
        
        # Call the function with the mocked object
        result = function_using_str_method(mock_obj)
        
        # Assert that the __str__ method of the object was called
        mock_obj.__str__.assert_called_once_with()
        
        # Assert that the result is the return value of the mocked __str__ method
        self.assertEqual(result, "Mocked __str__ implementation")

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

Best Practices for Unit Testing: When using unittest, patch, mock, and MagicMock, keep these best practices in mind:

Isolation: Isolate the code under test from external dependencies and focus on testing the specific functionality.
Readable and Maintainable: Write clear and concise test cases to make them easily understandable and maintainable.
Use Meaningful Names: Name your tests and mocks meaningfully to convey their purpose and intent.
Keep Tests Independent: Each test case should be independent and not rely on the state of previous tests.
Don’t Overuse Mocking: Mock only what’s necessary to achieve isolation; avoid excessive mocking, which can lead to brittle tests.

Real-World Use Cases: To illustrate the practical application of unittest, patch, mock, and MagicMock, let's look at two real-world examples.

Example 1: Testing a Database Connection: Suppose you have a function that interacts with a database. To test this function, you can use patch to mock the database connection and control its behaviour.

import unittest
from unittest.mock import patch
def fetch_data_from_db():
    # Connect to a database and fetch data
    pass
def process_data():
    data = fetch_data_from_db()
    # Process the data
class TestProcessData(unittest.TestCase):
    @patch('your_module.fetch_data_from_db')
    def test_process_data(self, mock_fetch_data_from_db):
        # Replace fetch_data_from_db with a mock object
        mock_fetch_data_from_db.return_value = "Mocked data"
        # Call the function under test
        result = process_data()
        # Assert that fetch_data_from_db was called
        mock_fetch_data_from_db.assert_called_once()
        # Perform other assertions based on the mocked data

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

Example 2: Testing a Snowflake Connection: Take a real-world example and mock the snowflake connection object that returns the data.

from unittest.mock import MagicMock,patch
import unittest
import snowflake.connector

def snowflake_connection(query) -> snowflake.connector.connection.SnowflakeConnection:
    sf_conn=snowflake.connector.connect(
        role='ACCOUNTADMIN',
        user="USER",
        password="ZXCD",
        account='xyx.eu-central-1',
        warehouse="COMPUTE_WH",
        schema='TEST',
        database="TEST",
    )
    sf_cursor = sf_conn.cursor()
    result=sf_cursor.execute(query)
    return result.fetchall()

class TestSnowflakeConnection(unittest.TestCase):

    @patch('snowflake.connector.connect')
    def test_snowflake_connection(self, mock_connect):
        # Create a MagicMock object for SnowflakeConnection
        mock_cursor = MagicMock()
        mock_conn = MagicMock()
        mock_conn.cursor.return_value = mock_cursor
        mock_connect.return_value = mock_conn

        # Set up test data
        query = 'SELECT * 1'

        # Call the function being tested
        result = snowflake_connection(query)

        # Assert that snowflake.connector.connect was called with the correct parameters
        mock_connect.assert_called_once_with(
            role='ACCOUNTADMIN',
            user="AYPAREEK",
            password="Timepass_90",
            account='xz10660.eu-central-1',
            warehouse="COMPUTE_WH",
            schema='TEST',
            database="TEST",
        )

        # Assert that the cursor's execute method was called with the query
        mock_cursor.execute.assert_called_once_with(query)

        # Assert that fetchall method was called on the result
        mock_result = mock_cursor.execute.return_value
        mock_result.fetchall.assert_called_once_with()

        # Assert that the result matches the return value of fetchall
        self.assertEqual(result, mock_result.fetchall.return_value)

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

Now wrap the snowflake connection into the udf and mock the udf.

from unittest.mock import MagicMock,patch
import unittest
import snowflake.connector

def snowflake_connection(query):
    sf_conn=snowflake.connector.connect(
        role='ACCOUNTADMIN',
        user="AYPAREEK",
        password="Timepass_90",
        account='xz10660.eu-central-1',
        warehouse="COMPUTE_WH",
        schema='TEST',
        database="TEST",
    )
    sf_cursor = sf_conn.cursor()
    return sf_cursor.execute(query)

def get_data(query):
    result=snowflake_connection(query)
    return result.fetchall()

class TestSnowflakeConnection(unittest.TestCase):

    @patch('<your_module>.snowflake_connection')
    def test_snowflake_connection(self, mock_data):
        # Replace fetch_data_from_db with a mock object
        mock_data.fetchall.return_value = 'mocked'
        result = get_data('select 1')
        mock_data.assert_called_once()    

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

Example 3: Testing a Boto3 method for accessing AWS resources. It’s time to patch the boto.client service and create a magic_mock object to replace it.

import unittest
from unittest.mock import patch, MagicMock
from typing import Dict
import json
import boto3

def fetch_secret(
    secret_key: str, secret_name: str = 'xyz/aws/database/variable'
) -> Dict[str, str]:
    # Create a Secrets Manager client using boto3
    secrets_client = boto3.client('secretsmanager', region_name='eu-central-1')
    # Retrieve the secret value from Secrets Manager using the provided secret name
    response = secrets_client.get_secret_value(SecretId=secret_name)
    # Extract the secret string from the response
    secret_value = response['SecretString']
    # Parse the secret string as JSON
    secret = json.loads(secret_value, strict=False)
    # Retrieve the specific secret using the provided key
    fetched_secret = secret[secret_key]
    # Return the fetched secret
    return fetched_secret

class FetchAwsSecret(unittest.TestCase):
    @patch('your_module.boto3.client')
    def test_fetch_secret(self, mock_client):
        # Replace requests.get with a mock object
        mock_secrets_client = MagicMock()
        mock_client.return_value = mock_secrets_client
        # Configure the mock to return a specific value for get_secret_value
        mock_secrets_client.get_secret_value.return_value = {
            'SecretString': '{"xyz": "secret1", "abc": "secret2"}'
        }
        # Set up test data
        secret_key = 'xyz'
        secret_name = 'secret1'
        # Call the function being tested
        actual_role_arn = fetch_secret(secret_key)
        # Assert the result
        self.assertEqual(actual_role_arn, secret_name)

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

Example 4: Testing API Calls: In another scenario, you should test a function that makes API calls. Using patch, you can isolate the API call and control its responses.

import unittest
from unittest.mock import patch, MagicMock
import requests
def call_api():
    response = requests.get("https://api.example.com/data")
    return response.json()
def process_api_data():
    data = call_api()
    # Process the API data
class TestProcessAPIData(unittest.TestCase):
    @patch('your_module.requests.get')
    def test_process_api_data(self, mock_requests_get):
        # Replace requests.get with a mock object
        mock_response = MagicMock()
        mock_response.json.return_value = {"key": "value"}
        mock_requests_get.return_value = mock_response
        # Call the function under test
        result = process_api_data()
        # Assert that requests.get was called with the correct URL
        mock_requests_get.assert_called_with("https://api.example.com/data")
        # Perform other assertions based on the mocked API response

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

Summary: In this article, we’ve explored the essential concepts of unittest, patch, mock, and MagicMock in Python unit testing. These tools enable you to isolate your code, control the behaviour of external dependencies, and create more robust and accurate unit tests.

Effective unit testing is a crucial practice for building reliable and maintainable software. It not only helps catch bugs early but also acts as documentation for your code, making it easier for other developers to understand and modify your code.

With a solid understanding of unittest, patch, mock, and MagicMock, you can significantly enhance the quality of your Python projects and make the development process smoother and more efficient. Happy testing!

If you think that this article is informative and helped you with what you are looking then give a clap and follow my medium account( datageeks.medium.com ) and feel free to write in the comments if you have any doubts about this topic.

By signing up as a member (https://datageeks.medium.com/membership), you can read every story and help the authors on Medium.