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

user_8942319 0
Viewed 4301

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?

boto3 amazon-s3 mocking unit-testing python
4 Answers

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.dump(body_json).encode("utf-8")

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

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

mock_get_object.return_value = mocked_response
user_76996760

The below code worked for me. referred answer: https://stackoverflow.com/a/64642433/12385686

import json
from botocore.response import StreamingBody
import io
    
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()

body = StreamingBody(
    io.BytesIO(body_encoded),
    len(body_encoded)
)

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

mock_get_object.return_value = mocked_response
user_123856860

works for me I used like

@patch('src.handler.boto3.client')
def test_AccountIDs(self, client:MagicMock):
    client.return_value = s3_client
    
    body_encoded = open('accounts.csv').read().encode()
    mock_stream = StreamingBody(io.BytesIO(body_encoded),len(body_encoded))
    s3_stubber.add_response('get_object', { 'Body' : mock_stream})

    with s3_stubber:
        res = handler.getAccountIDs()
        self.assertListEqual(res,['one', 'two', 'three'])

thanks for the solution !!!! :)

user_73319730

In my case, I was trying to mock a .json file as response

# my_file.json
{
    "key1": "value_1",
    "key2": "value_2",
    "key3": "value_3"
}

Using BytesIO instead of StringIO in @YisusThreepwood answer works

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

Otherwise you get

TypeError: initial_value must be str or None, not bytes.

user_124833460
Related