如何模拟 boto3 的 StreamingBody 对象，以便在 Python 中使用 BytesIO 进行处理？

Question

我正在对一个函数进行单元测试，该函数将元素从 S3 对象转换为熊猫数据帧，并且需要模拟从 boto3 返回的 StreamingBody 对象

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

此处记录了self.get_object(key_name(的响应

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

所以我需要模拟那个 StreamingBody(( 对象并让我的模拟函数返回它。

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)

但我遇到了TypeError: 'StreamingBody' object is not subscriptable

有什么提示吗？

Answer 1

S3 客户端返回一个字典，而您模拟的 S3 客户端返回一个 StreamingBody。您模拟的 S3 客户端应返回类似

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

Answer 2

下面的代码对我有用。 参考答案：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

Answer 3

对我有用，我用过

@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'])

感谢您的解决方案!!! :)

Answer 4

就我而言，我试图模拟一个.json文件作为响应

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

在答案中使用BytesIO而不是StringIO@YisusThreepwood有效

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

否则你会得到

类型错误：initial_value必须是 str 或 None，而不是字节。