From 757d03ebc00b6b97e74aac4bee58fd572662b47b Mon Sep 17 00:00:00 2001
From: Ilya Pekelny <pekelny@gmail.com>
Date: Mon, 23 Jul 2018 18:39:34 +0300
Subject: [PATCH] Record empty body if an object is of bytes type

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.
---
 placebo/pill.py       | 19 ++++++++++++++++++-
 placebo/serializer.py | 14 ++++++++++----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/placebo/pill.py b/placebo/pill.py
index 1589da5..889ffdd 100644
--- a/placebo/pill.py
+++ b/placebo/pill.py
@@ -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
 
@@ -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):
         """
@@ -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)
 
diff --git a/placebo/serializer.py b/placebo/serializer.py
index cb18e99..09df348 100644
--- a/placebo/serializer.py
+++ b/placebo/serializer.py
@@ -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"""
 
@@ -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
@@ -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")