How do you change the output of mocked methods, using the Mock framework?

Question

I have a module like the following:

'''example.py'''
from some.module import Session

def init_system():
    var_one = Session.query(SQLAModelA).all()
    var_two = [(x.attr_B_a, x) for x in Session.query(SQLAModelB).all()]

    for value in var_one:
        something = value.attr_A_a + value.attr_A_c
    # Will do something with var_two etc (which has different attributes)
    something_else = [x.attr_B * 2 for x in var_two]
    return result

Session refers to an SQLAlchemy session handler which talks to a DB. For testing purposes I need to be able to Mock the query method in Session object (or maybe just the Session object?) so I can test init_system without it actually interacting with the database. How do I use the Mock framework to do this?

Really the bit that puzzles me is, if I mock the query() method how do I change the output of it's all method dependent on the class that is passed to query() ??

I'm hoping I can do something along the lines of:

def test_init_numbers():
    import
    my_mock_object = MagicMock()
    with patch('some.module.Session.query', my_mock_object):
        result = example.init_system()
        assert result == ['expected', 'result']

However, I don't know what I need to do to my_mock_object to make it mimic query() correctly. I guess it needs to implement the all() method which returns an iterator. And, I know that I could create some more mock objects to go in that returned iterator, making sure that they have values for the appropriate attributes, eg:

var_one_element_mock = MagicMock(spec=SQLAModelA)
var_one_element_mock.attr_A_a = 12

var_two_element_mock = MagicMock(spec=SQLAModelB)
var_one_element_mock.attr_B = 100

And, of course, I can create iterables of different versions of those MagicMock objects which can be used by the init_system function, eg, in the for value in var_one line.

But I don't know how to piece all the mocked objects together so I can just mock out the Session.query method taking into account that a class is passed to it, which affects the contents of the iterable it returns.

Thanks!

Answer 1

I don't know if you still need help with this, but if you want to control what happens when a MagicMock object is called, you can pass a function as its side_effect argument. Something like this should work.

def test_init_numbers():
    import
    my_mock_object = MagicMock()
    def return_iterator():
      ## create and return your iterator here

    my_mock_object.all = MagicMock(side_effect=return_iterator)
    with patch('some.module.Session.query', my_mock_object):
        result = example.init_system()
        assert result == ['expected', 'result']