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?

Did you get the solution? I'm exactly facing the same problem — Tula, Sep 09 '20 at 07:05
@Tula, No I didn't figure it out. I took another approach - figuring those aws/pandas functions are themselves maintained and tested and there's not a lot of value to be had by implementing that granular of a unit test. — sam, Sep 14 '20 at 14:59

score 8 · Answer 1 · edited Dec 21 '22 at 11:04

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

`json.dump` won't work, I think you mean `json.dumps`. this solution did not work for me, Akshay kamath B did — Ralph Willgoss, Jan 11 '22 at 14:08

score 5 · Answer 2 · answered Mar 09 '21 at 22:02

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

score 0 · Answer 3 · answered Apr 19 '22 at 21:34

0

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 !!!! :)

answered Apr 19 '22 at 21:34

Gerardo Lucio

1

Your answer could be improved with additional supporting information. Please [edit] to add further details, such as citations or documentation, so that others can confirm that your answer is correct. You can find more information on how to write good answers [in the help center](/help/how-to-answer). – Community Apr 19 '22 at 23:06

score 0 · Answer 4 · answered Aug 31 '22 at 00:09

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.

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

4 Answers4