How do I mock boto3's StreamingBody object for processing with BytesIO in Python?

Question

I'm unittesting a function that transforms an element from an S3 object into a pandas DataFrame and need to mock the returned StreamingBody object from boto3

file.py

def object_to_df(self, key_name, dtypes):
    s3_object = self.get_object(key_name=key_name)
    if s3_object is not None:
        object_df = pandas.read_csv(
            io.BytesIO(s3_object["Body"].read()), dtype=dtypes
        )
        return object_df

The response of self.get_object(key_name) is documented here

{
    'Body': StreamingBody(),
    'DeleteMarker': True|False,
    'AcceptRanges': 'string',
    ...
}

So I need to mock that StreamingBody() object and have my mock function return that.

test.py

import unittest
import pandas
from io import StringIO
from unittest.mock import patch, Mock
from path.to.file import custom_class
from botocore.response import StreamingBody

class TestS3Class(unittest.TestCase):
    """TestCase for path_to/file.py"""

    def setUp(self):
        """Creates an instance of the live class for testing"""
        self.s3_test_client = S3()


    @patch('path.to.class.get_object')
    def test_object_to_df(self, mock_get_object):
        """"""
        mock_response = {'Body': [{'Candidate': 'Black Panther', 'Votes': 3},
                        {'Candidate': 'Captain America: Civil War', 'Votes': 8},
                        {'Candidate': 'Guardians of the Galaxy', 'Votes': 8},
                        {'Candidate': "Thor: Ragnarok", 'Votes': 1}
                    ]}
        mock_stream = StreamingBody(StringIO(str(mock_response)), len(str(mock_response)))
        mock_get_object.return_value = mock_stream
        self.assertIsInstance(self.s3_test_client.object_to_df(key_name='key_name', dtypes=str), pandas.DataFrame)

But I'm running into TypeError: 'StreamingBody' object is not subscriptable

Any hints?

Answer 1

The S3 client returns a dict and your mocked S3 client is returning a StreamingBody. Your mocked S3 client should return something like

body_json = {
    'Body': [
        {'Candidate': 'Black Panther', 'Votes': 3},
        {'Candidate': 'Captain America: Civil War', 'Votes': 8},
        {'Candidate': 'Guardians of the Galaxy', 'Votes': 8},
        {'Candidate': "Thor: Ragnarok", 'Votes': 1},
    ]
}

body_encoded = json.dumps(body_json).encode('utf-8')

body = StreamingBody(
    StringIO(body_encoded),
    len(body_encoded)
)

mocked_response = {
    'Body': body,
    ...
}

mock_get_object.return_value = mocked_response