Mock boto3 batch write for unit testing in python

Question

I would like to test the batch_write method using unit test.

class DataService:
    def __init__(self, table):
        dynamodb = boto3.resource('dynamodb', region_name='us-west-2')
        self.db = dynamodb.Table(self.table)


    def batch_write(self, items):
        with self.db.batch_writer() as batch:
            for item in items:
                batch.put_item(Item=item)

I want to mock the batch_writer(). I have seen usage of mock_open used to mock file open methods. But I don't really understand can i mock this streaming batch_writer() call.

Answer 1

I came across this issue as well this is how i solved it.

    from mock import MagicMock
    class TestClass():
        def __enter__(self, *args):
            return self
        def __exit__(self, *args):
            pass

    cls = TestClass()
    cls.put_item = MagicMock()

    repository.table.batch_writer = MagicMock()
    repository.table.batch_writer.return_value = cls

    print cls.put_item.call_args_list

Now i can pass my mock test class in place of the batch_writer to mock calls or check params. I know this is an old question but it is one that I struggled with and wished someone would have posted a code example if this is easy to do. If there is better easier way using patch that anyone has please post.

Answer 2

I wrote a slightly different version, based on Codyj110's idea, but it won't require to create the TestClass.

Instead of creating the new class, I'm just setting the __enter__ and __exit__ methods with mocked values.

class TestClassRequiringMockOnDynamoDbBatchWriter:
    @patch('boto3.resource')
    def test_mocked_batch_writer(self, mock_dynamo_db):
        mock_db, mock_batch_writer = self._get_mocked_dynamo_objects(fake_exception)
        mock_dynamo_db.return_value = mock_db

        assert mock_batch_writer.put_item.call_count == my_expected_calls


    @staticmethod
    def _get_mocked_dynamo_objects(expected_value):
        mock_batch_writer = Mock()
        mock_batch_writer.__enter__ = Mock(return_value=mock_batch_writer)
        mock_batch_writer.__exit__ = Mock(return_value=None)

        # Use side_effect or return_value according to your intention
        mock_batch_writer.put_item.side_effect = expected_value

        mock_table = Mock()
        mock_table.batch_writer.return_value = mock_batch_writer

        mock_db = Mock()
        mock_db.Table.return_value = mock_table

        return mock_db, mock_batch_writer

I hope it helps someone!