Python how to mock a function that's returned by a higher-order function

Question

I'd like to mock the Pandas function read_csv, which in my code is returned dynamically by a higher-order function, but it seems the pytest patch doesn't really allow it.

import pandas as pd

format_read_func_mapping = {"csv": pd.read_csv, "parquet": pd.read_parquet}

def my_func(s3_path, file_format):
    read_func = format_read_func_mapping[file_format]
    df = read_func(f"{s3_path}")
    return df

When I run my unit test, I have the following

from unittest.mock import patch

@patch(f"{MODULE_PATH}.pd.read_csv")
def test_my_func(self, mock_read_csv):
    mock_read_csv.side_effect = [my_test_data_frame]
    my_func(dummy_s3_path, "csv")

I expect the test would use the mock and return my test data, but it doesn't, it actually called pandas.read_csv to try to read from dummy_s3_path.

However, if I don't dynamically return pd.read_csv from a mapping, just directly use pd.read_csv(f"{s3_path}") in my_func, the mock works well.

Is this a limit of pytest? Or the way I mock is incorrect?

Thanks.

Answer 1

The problem is that you're defining the format_read_func_mapping dict at the module level, so when the module is imported the dict is created with the value of the key 'csv' being a reference to the original pd.read_csv, which does not get modified when you then patch pd.read_csv afterwards.

To patch format_read_func_mapping in the global namespace of the module, you can use patch.dict to patch the specific key of the dict instead:

import module

@patch.dict(module.format_read_func_mapping, {"csv": Mock(side_effect=[my_test_data_frame])})
def test_my_func():
    my_func(dummy_s3_path, "csv")