在pytest中的函数内为同一调用返回不同的模拟结果

Question

I have a test function where I successfully patched a function call to return a mock result. It works for the one time the patched function is called:

@patch(‘db.get_session’)
def test_some_function(mock_session):
    mock_session.return_value.query.return_value = MockResult(1)
    assert some_function()

However, I'm trying to test a function that uses get_session().query() multiple times and should return a different result each time. I'm not sure how to mock this.

@patch(‘db.get_session’)
def test_some_other_function(mock_session):
    # First call should return one
    mock_session.return_value.query.return_value = MockResult(1)
    # Second call should return no results:
    mock_session.return_value.query.return_value = MockResult(0)
    assert some_other_function()

def some_other_function():
    session = get_session()
    one = session.query('some_query')
    zero = session.query('some_other_query')
    if one == 1 and zero == 0:
        return True

I've taken a look at side_effect but I'm not sure how to set it up since I'm not changing the straight return_value but rather the return_value of the return_value (if that makes sense).

Answer 1

the side_effect attribute of a mock can contain an iterable -- each item in the iterable will be returned in order:

>>> m = mock.Mock(side_effect=(1, 2))
>>> m()
1
>>> m()
2

combining this with your example:

    mock_session.return_value.query.side_effect = (MockResult(1), MockResult(0))