Search code examples
pythonunit-testingpython-mock

Mocking return value of a nested call in Python mock library


Brand new to this library

Here is the call stack of my mocked object

[call(),
 call('test'),
 call().instance('test'),
 call().instance().database('test'),
 call().instance().database().snapshot(),
 call().instance().database().snapshot().__enter__(),
 call().instance().database().snapshot().__enter__().execute_sql('SELECT * FROM users'),
 call().instance().database().snapshot().__exit__(None, None, None),
 call().instance().database().snapshot().__enter__().execute_sql().__iter__()]

Here is the code I have used

    @mock.patch('testmodule.Client')
    def test_read_with_query(self, mock_client):
        mock = mock_client()
        pipeline = TestPipeline()
        records = pipeline | ReadFromSpanner(TEST_PROJECT_ID, TEST_INSTANCE_ID, self.database_id).with_query('SELECT * FROM users')
        pipeline.run()
        print mock_client.mock_calls
        exit()

I want to mock this whole stack that eventually it gives me some fake data which I will provide as a return value.

The code being tested is


        spanner_client = Client(self.project_id)
        instance = spanner_client.instance(self.instance_id)
        database = instance.database(self.database_id)

        with database.snapshot() as snapshot:
            results = snapshot.execute_sql(self.query)

So my requirements is that the results variable should contain the data I will provide.

How can I provide a return value to such a nested calls

Thanks


Solution

  • Create separate MagicMock instances for the instance, database and snapshot objects in the code under test. Use return_value to configure the return values of each method. Here is an example. I simplified the method under test to just be a free standing function called mut.

    # test_module.py : the module under test
    class Client:
        pass
    
    def mut(project_id, instance_id, database_id, query):
        spanner_client = Client(project_id)
        instance = spanner_client.instance(instance_id)
        database = instance.database(database_id)
    
        with database.snapshot() as snapshot:
            results = snapshot.execute_sql(query)
            return results
    
    # test code (pytest)
    from unittest.mock import MagicMock
    from unittest import mock
    
    from test_module import mut
    
    @mock.patch('test_module.Client')
    def test_read_with_query(mock_client_class):
        mock_client = MagicMock()
        mock_instance = MagicMock()
        mock_database = MagicMock()
        mock_snapshot = MagicMock()
        expected = 'fake query results'
    
        mock_client_class.return_value = mock_client
        mock_client.instance.return_value = mock_instance
        mock_instance.database.return_value = mock_database
        mock_database.snapshot.return_value = mock_snapshot
        mock_snapshot.execute_sql.return_value = expected
        mock_snapshot.__enter__.return_value = mock_snapshot
    
        observed = mut(29, 42, 77, 'select *')
    
        mock_client_class.assert_called_once_with(29)
        mock_client.instance.assert_called_once_with(42)
        mock_instance.database.assert_called_once_with(77)
        mock_database.snapshot.assert_called_once_with()
        mock_snapshot.__enter__.assert_called_once_with()
        mock_snapshot.execute_sql.assert_called_once_with('select *')
        assert observed == expected
    

    This test is kind of portly. Consider breaking it apart by using a fixture and a before function that sets up the mocks.