How to mock chained function calls in python?

Question

I'm using the mock library written by Michael Foord to help with my testing on a django application.

I'd like to test that I'm setting up my query properly, but I don't think I need to actually hit the database, so I'm trying to mock out the query.

I can mock out the first part of the query just fine, but I am not getting the results I'd like when I chain additional things on.

The function:

    @staticmethod
    def get_policies(policy_holder, current_user):
        if current_user.agency:
            return Policy.objects.filter(policy_holder=policy_holder, version__agency=current_user.agency).distinct()
        else:
            return Policy.objects.filter(policy_holder=policy_holder)

and my test: The first assertion passes, the second one fails.

    def should_get_policies_for_agent__user(self):
        with mock.patch.object(policy_models.Policy, "objects") as query_mock:
            user_mock = mock.Mock()
            user_mock.agency = "1234"
            policy_models.Policy.get_policies("policy_holder", user_mock)
            self.assertEqual(query_mock.method_calls, [("filter", (), {
                'policy_holder': "policy_holder",
                'version__agency': user_mock.agency,
            })])
            self.assertTrue(query_mock.distinct.called)

I'm pretty sure the issue is that the initial query_mock is returning a new mock after the .filter() is called, but I don't know how to capture that new mock and make sure .distinct() was called on it.

Is there a better way to be testing what I am trying to get at? I'm trying to make sure that the proper query is being called.

Answer 1

Each mock object holds onto the mock object that it returned when it is called. You can get a hold of it using your mock object's return_value property.

For your example,

self.assertTrue(query_mock.distinct.called)

distinct wasn't called on your mock, it was called on the return value of the filter method of your mock, so you can assert that distinct was called by doing this:

self.assertTrue(query_mock.filter.return_value.distinct.called)