Skip to content

Commit

Permalink
Record empty body if an object is of bytes type
Browse files Browse the repository at this point in the history
During recording a S3 object retrieving request the library raises
`TypeError` while the response is totally valid. Response json data
 is copied including streaming body to handle potential
issues with binary data without affecting original binary payload.
  • Loading branch information
I159 committed Jul 23, 2018
1 parent fba32ff commit 757d03e
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 5 deletions.
19 changes: 18 additions & 1 deletion placebo/pill.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import json
import io
import os
import glob
import re
import shutil
import uuid
import logging

Expand Down Expand Up @@ -225,6 +228,19 @@ def get_next_file_path(self, service, operation):
raise IOError('response file ({0}) not found'.format(fn))
return fn

def copy_json_data(self, json_data):
try:
return copy.deepcopy(json_data)
except TypeError:
copy_json_data = {k:v for k, v in json_data.items() if k != 'Body'}
buffer = io.BytesIO()
shutil.copyfileobj(json_data['Body'], buffer)
buffer.seek(0)
json_data['Body']._amount_read = 0
json_data['Body']._raw_stream = io.BytesIO(buffer.getvalue())
copy_json_data['Body'] = buffer
return copy_json_data

def save_response(self, service, operation, response_data,
http_response=200):
"""
Expand All @@ -240,7 +256,8 @@ def save_response(self, service, operation, response_data,
filepath = self.get_new_file_path(service, operation)
LOG.debug('save_response: path=%s', filepath)
json_data = {'status_code': http_response,
'data': response_data}
'data': self.copy_json_data(response_data)}

with open(filepath, 'w') as fp:
json.dump(json_data, fp, indent=4, default=serialize)

Expand Down
14 changes: 10 additions & 4 deletions placebo/serializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,10 @@
# limitations under the License.

from datetime import datetime, timedelta, tzinfo
from botocore.response import StreamingBody
import io
from six import StringIO


class UTC(tzinfo):
"""UTC"""

Expand All @@ -28,8 +29,10 @@ def tzname(self, dt):
def dst(self, dt):
return timedelta(0)


utc = UTC()


def deserialize(obj):
"""Convert JSON dicts back into objects."""
# Be careful of shallow copy here
Expand Down Expand Up @@ -66,10 +69,13 @@ def serialize(obj):
result['second'] = obj.second
result['microsecond'] = obj.microsecond
return result
if isinstance(obj, StreamingBody):
if isinstance(obj, io.BytesIO):
result['body'] = obj.read()
obj._raw_stream = StringIO(result['body'])
obj._amount_read = 0
try:
result['body'] = result['body'].decode('utf-8')
except UnicodeError:
# Could be turned back to bytes `bytes(result['body'])`
result['body'] = list(result['body'])
return result
# Raise a TypeError if the object isn't recognized
raise TypeError("Type not serializable")

0 comments on commit 757d03e

Please sign in to comment.