Over the last few days, I’ve been tackling a problem related to mocking a function that accepts parameters. More specifically, I’ve been trying to test the following logic:
def my_function(df) -> pd.DataFrame:
return (
df.groupby(self.groupby_columns)[
[column_a, column_b]
]
.apply(my_function_that_call_external_api)
.reset_index()
)
def my_function_that_call_external_api(imput_df: pd.DataFrameGroupBy) -> pd.DataFrame:
df = preprocess(df)
predictions = external_api_call(df)
return postprocess_predictions(predictions)
For the sake of this discussion, consider external_api_call
as a function that calls Prophet, both for fitting and predicting.
To test this functionality, I need to mock the external_api_call
. Here’s my attempt:
from unittest.mock import patch
from functools import partial
def generate_mock_reposnse(self, column_a_modifier: int, column_b_modifie: int, input_df: pd.DataFrame) -> pd.DataFrame:
"""Generate a dummy external api response with random prediction.
"""
response = pd.DataFrame(
{
"ds":np.full_like(input_df.ds, column_a_modifier, dtype=float),
"yhat": np.full_like(input_df.ds, column_b_modifie, dtype=float),
}
)
return response
@patch("external_api_call")
def test_my_function_with_mock(self, api_call_mock):
df = self.generate_df()
api_call_mock.side_effect = [partial(generate_mock_reposnse,
column_a_modifier = a,
column_b_modifie = b)
for a in a_values
for b in b_values
)]
result = my_function(df)
self.validate_result(result)
Unfortunately, this code does not work as expected because the partial function is never called. I propose modifying the behavior of MagicMock.side_effect to enable the execution of callable functions.