How to add args for mocked functions with MagicMock

Question

I'm using MagicMock from unittest to mock an object in a class for unit tests. I am able to use the return_value to set return values for function calls but want to make it dynamic so that the response changes based on the args passed.

Here's the code for the original function(_session is created from sqlalchemy sessionmaker) :

def table_filter(self, filter_exp):
    result = self._session.query(Table1).filter(text(filter_exp)).all()
    return result

Code for unit test :

def test_filter(mock_object):

    db_connection = DBConnection("", "", 0, "", "")
    mock_session = MagicMock()
    attrs = {
            "query.return_value.filter.return_value.all.return_value": [mock_object]
        }
    mock_session.configure_mock(**attrs)
    setattr(db_connection, "_session", mock_session)

    assert db_connection.table_filter("WHERE x = 1") == [mock_object]

I am just overwriting the session attribute of the db_connection object and replacing it with a mocked version. How can i update the attrs so that args can be defined, something like this:

attrs = {
            "query.return_value.filter("WHERE x = 1").return_value.all.return_value": [mock_object_1]
            "query.return_value.filter("WHERE x = 2").return_value.all.return_value": [mock_object_2]
        }

I'm using sqlalchemy and want to mock the underlying session object created using sqlalchemy's sessionmaker. So that when unit test are run, all the normal code runs and the interactions with the database using the session object is mocked.

It is possible to use side_effect to mock the outer functions return value but doing it on the lower function call seems tricky

How to add args for mocked functions with MagicMock

0 Answers0