Merge pull request #38 from yugangw-msft/bin2

Support playback with binary and large response payload

Author: tjprescott

requirements.txt

@@ -1,6 +1,6 @@
 codecov==2.0.9
 mock==2.0.0;python_version<="2.7"
 nose==1.3.7
-pylint==1.7.1
+pylint==1.8.2
 
 -e .

src/azure_devtools/scenario_tests/__init__.py

@@ -7,7 +7,7 @@
 from .exceptions import AzureTestError
 from .decorators import live_only, record_only
 from .patches import mock_in_unit_test, patch_time_sleep_api, patch_long_run_operation_delay
-from .preparers import AbstractPreparer, SingleValueReplacer
+from .preparers import AbstractPreparer, SingleValueReplacer, AllowLargeResponse
 from .recording_processors import (
     RecordingProcessor, SubscriptionRecordingProcessor,
     LargeRequestBodyProcessor, LargeResponseBodyProcessor, LargeResponseBodyReplacer,
@@ -18,7 +18,7 @@
 __all__ = ['IntegrationTestBase', 'ReplayableTest', 'LiveTest',
            'AzureTestError',
            'mock_in_unit_test', 'patch_time_sleep_api', 'patch_long_run_operation_delay',
-           'AbstractPreparer', 'SingleValueReplacer',
+           'AbstractPreparer', 'SingleValueReplacer', 'AllowLargeResponse',
            'RecordingProcessor', 'SubscriptionRecordingProcessor',
            'LargeRequestBodyProcessor', 'LargeResponseBodyProcessor', 'LargeResponseBodyReplacer',
            'OAuthRequestResponsesFilter', 'DeploymentNameReplacer', 'GeneralNameReplacer',

src/azure_devtools/scenario_tests/base.py

@@ -162,6 +162,7 @@ def _process_request_recording(self, request):
         return request
 
     def _process_response_recording(self, response):
+        from .utilities import is_text_payload
         if self.in_recording:
             # make header name lower case and filter unwanted headers
             headers = {}
@@ -171,7 +172,7 @@ def _process_response_recording(self, response):
             response['headers'] = headers
 
             body = response['body']['string']
-            if body and not isinstance(body, six.string_types):
+            if is_text_payload(response) and body and not isinstance(body, six.string_types):
                 response['body']['string'] = body.decode('utf-8')
 
             for processor in self.recording_processors:

src/azure_devtools/scenario_tests/preparers.py

@@ -8,7 +8,7 @@
 import functools
 
 from .base import ReplayableTest
-from .utilities import create_random_name
+from .utilities import create_random_name, is_text_payload
 from .recording_processors import RecordingProcessor
 
 
@@ -111,15 +111,15 @@ def process_request(self, request):
             request.uri = request.uri.replace(quote_plus(self.random_name),
                                               quote_plus(self.moniker))
 
-        if request.body:
+        if is_text_payload(request) and request.body:
             body = str(request.body)
             if self.random_name in body:
                 request.body = body.replace(self.random_name, self.moniker)
 
         return request
 
     def process_response(self, response):
-        if response['body']['string']:
+        if is_text_payload(response) and response['body']['string']:
             response['body']['string'] = response['body']['string'].replace(self.random_name,
                                                                             self.moniker)
 
@@ -129,6 +129,20 @@ def process_response(self, response):
         return response
 
 
+# Function wise, enabling large payload recording has nothing to do with resource preparers
+# We still base on it so that this decorator can chain with other preparers w/o too much hacks
+class AllowLargeResponse(AbstractPreparer):
+    def __init__(self, size_kb=1024):
+        self.size_kb = size_kb
+        super(AllowLargeResponse, self).__init__('nanana', 20)
+
+    def create_resource(self, _, **kwargs):
+        from azure_devtools.scenario_tests import LargeResponseBodyProcessor
+        large_resp_body = next((r for r in self.test_class_instance.recording_processors
+                                if isinstance(r, LargeResponseBodyProcessor)), None)
+        if large_resp_body:
+            large_resp_body._max_response_body = self.size_kb  # pylint: disable=protected-access
+
 # Utility
 
 def is_preparer_func(fn):

src/azure_devtools/scenario_tests/recording_processors.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from .utilities import is_text_payload, is_json_payload
 
 class RecordingProcessor(object):
     def process_request(self, request):  # pylint: disable=no-self-use
@@ -31,13 +32,13 @@ def __init__(self, replacement):
     def process_request(self, request):
         request.uri = self._replace_subscription_id(request.uri)
 
-        if request.body:
+        if is_text_payload(request) and request.body:
             request.body = self._replace_subscription_id(request.body.decode()).encode()
 
         return request
 
     def process_response(self, response):
-        if response['body']['string']:
+        if is_text_payload(response) and response['body']['string']:
             response['body']['string'] = self._replace_subscription_id(response['body']['string'])
 
         self.replace_header_fn(response, 'location', self._replace_subscription_id)
@@ -66,7 +67,7 @@ def __init__(self, max_request_body=128):
         self._max_request_body = max_request_body
 
     def process_request(self, request):
-        if request.body and len(request.body) > self._max_request_body * 1024:
+        if is_text_payload(request) and request.body and len(request.body) > self._max_request_body * 1024:
             request.body = '!!! The request body has been omitted from the recording because its ' \
                            'size {} is larger than {}KB. !!!'.format(len(request.body),
                                                                      self._max_request_body)
@@ -81,34 +82,42 @@ def __init__(self, max_response_body=128):
         self._max_response_body = max_response_body
 
     def process_response(self, response):
-        length = len(response['body']['string'] or '')
-        if length > self._max_response_body * 1024:
-            response['body']['string'] = \
-                "!!! The response body has been omitted from the recording because it is larger " \
-                "than {} KB. It will be replaced with blank content of {} bytes while replay. " \
-                "{}{}".format(self._max_response_body, length, self.control_flag, length)
-
+        if is_text_payload(response):
+            length = len(response['body']['string'] or '')
+            if length > self._max_response_body * 1024:
+
+                if is_json_payload(response):
+                    from .preparers import AllowLargeResponse  # pylint: disable=cyclic-import
+                    raise ValueError("The json response body exceeds the default limit of {}kb. Use '@{}' "
+                                     "on your test method to increase the limit or update test logics to avoid "
+                                     "big payloads".format(self._max_response_body, AllowLargeResponse.__name__))
+
+                response['body']['string'] = \
+                    "!!! The response body has been omitted from the recording because it is larger " \
+                    "than {} KB. It will be replaced with blank content of {} bytes while replay. " \
+                    "{}{}".format(self._max_response_body, length, self.control_flag, length)
         return response
 
 
 class LargeResponseBodyReplacer(RecordingProcessor):
     def process_response(self, response):
-        import six
-        body = response['body']['string']
+        if is_text_payload(response) and not is_json_payload(response):
+            import six
+            body = response['body']['string']
 
-        # backward compatibility. under 2.7 response body is unicode, under 3.5 response body is
-        # bytes. when set the value back, the same type must be used.
-        body_is_string = isinstance(body, six.string_types)
+            # backward compatibility. under 2.7 response body is unicode, under 3.5 response body is
+            # bytes. when set the value back, the same type must be used.
+            body_is_string = isinstance(body, six.string_types)
 
-        content_in_string = (response['body']['string'] or b'').decode('utf-8')
-        index = content_in_string.find(LargeResponseBodyProcessor.control_flag)
+            content_in_string = (response['body']['string'] or b'').decode('utf-8')
+            index = content_in_string.find(LargeResponseBodyProcessor.control_flag)
 
-        if index > -1:
-            length = int(content_in_string[index + len(LargeResponseBodyProcessor.control_flag):])
-            if body_is_string:
-                response['body']['string'] = '0' * length
-            else:
-                response['body']['string'] = bytes([0] * length)
+            if index > -1:
+                length = int(content_in_string[index + len(LargeResponseBodyProcessor.control_flag):])
+                if body_is_string:
+                    response['body']['string'] = '0' * length
+                else:
+                    response['body']['string'] = bytes([0] * length)
 
         return response
 
@@ -122,6 +131,7 @@ def process_request(self, request):
         import re
         if not re.match('https://login.microsoftonline.com/([^/]+)/oauth2/token', request.uri):
             return request
+        return None
 
 
 class DeploymentNameReplacer(RecordingProcessor):
@@ -159,7 +169,7 @@ def process_request(self, request):
         for old, new in self.names_name:
             request.uri = request.uri.replace(old, new)
 
-            if request.body:
+            if is_text_payload(request) and request.body:
                 body = str(request.body)
                 if old in body:
                     request.body = body.replace(old, new)
@@ -168,7 +178,7 @@ def process_request(self, request):
 
     def process_response(self, response):
         for old, new in self.names_name:
-            if response['body']['string']:
+            if is_text_payload(response) and response['body']['string']:
                 response['body']['string'] = response['body']['string'].replace(old, new)
 
             self.replace_header(response, 'location', old, new)

src/azure_devtools/scenario_tests/tests/test_recording_processor.py

@@ -78,6 +78,7 @@ def test_subscription_recording_processor_for_request(self):
             mock_request = mock.Mock()
             mock_request.uri = template.format(mock_sub_id)
             mock_request.body = self._mock_subscription_request_body(mock_sub_id)
+            mock_request.headers = {'content-type': 'application/json'}
 
             rp.process_request(mock_request)
             self.assertEqual(mock_request.uri, template.format(replaced_subscription_id))
@@ -106,11 +107,35 @@ def test_subscription_recording_processor_for_response(self):
             mock_response = dict({'body': {}})
             mock_response['body']['string'] = template.format(mock_sub_id)
             mock_response['headers'] = {'Location': [location_header_template.format(mock_sub_id)],
-                                        'azure-asyncoperation': [asyncoperation_header_template.format(mock_sub_id)]}
+                                        'azure-asyncoperation': [asyncoperation_header_template.format(mock_sub_id)],
+                                        'content-type': ['application/json']}
             rp.process_response(mock_response)
             self.assertEqual(mock_response['body']['string'], template.format(replaced_subscription_id))
 
             self.assertSequenceEqual(mock_response['headers']['Location'],
                                      [location_header_template.format(replaced_subscription_id)])
             self.assertSequenceEqual(mock_response['headers']['azure-asyncoperation'],
                                      [asyncoperation_header_template.format(replaced_subscription_id)])
+
+
+    def test_recording_processor_skip_body_on_unrecognized_content_type(self):
+        location_header_template = 'https://graph.windows.net/{}/directoryObjects/' \
+                                   'f604c53a-aa21-44d5-a41f-c1ef0b5304bd/Microsoft.DirectoryServices.Application'
+        replaced_subscription_id = str(uuid.uuid4())
+        rp = SubscriptionRecordingProcessor(replaced_subscription_id)
+
+        mock_sub_id = str(uuid.uuid4())
+        mock_response = dict({'body': {}})
+        mock_response['body']['string'] = mock_sub_id
+        mock_response['headers'] = {
+            'Location': [location_header_template.format(mock_sub_id)],
+            'content-type': ['application/foo']
+        }
+
+        # action
+        rp.process_response(mock_response)
+
+        # assert
+        self.assertEqual(mock_response['body']['string'], mock_sub_id)  # body unchanged
+        self.assertEqual(mock_response['headers']['Location'],
+                         [location_header_template.format(replaced_subscription_id)])  # header changed

src/azure_devtools/scenario_tests/tests/test_utilities.py

@@ -4,7 +4,9 @@
 # --------------------------------------------------------------------------------------------
 
 import unittest
-from azure_devtools.scenario_tests.utilities import create_random_name, get_sha1_hash
+import mock
+from azure_devtools.scenario_tests.utilities import (create_random_name, get_sha1_hash,
+                                                     is_text_payload, is_json_payload)
 
 
 class TestUtilityFunctions(unittest.TestCase):
@@ -80,3 +82,37 @@ def test_get_sha1_hash(self):
             f.seek(0)
             hash_value = get_sha1_hash(f.name)
             self.assertEqual('6487bbdbd848686338d729e6076da1a795d1ae747642bf906469c6ccd9e642f9', hash_value)
+
+    def test_text_payload(self):
+        http_entity = mock.MagicMock()
+        headers = {}
+        http_entity.headers = headers
+
+        headers['content-type'] = 'foo/'
+        self.assertFalse(is_text_payload(http_entity))
+
+        headers['content-type'] = 'text/html; charset=utf-8'
+        self.assertTrue(is_text_payload(http_entity))
+
+        headers['content-type'] = 'APPLICATION/JSON; charset=utf-8'
+        self.assertTrue(is_text_payload(http_entity))
+
+        headers['content-type'] = 'APPLICATION/xml'
+        self.assertTrue(is_text_payload(http_entity))
+
+        http_entity.headers = None  # default to text mode if there is no header
+        self.assertTrue(is_text_payload(http_entity))
+
+    def test_json_payload(self):
+        http_entity = mock.MagicMock()
+        headers = {}
+        http_entity.headers = headers
+
+        headers['content-type'] = 'APPLICATION/JSON; charset=utf-8'
+        self.assertTrue(is_json_payload(http_entity))
+
+        headers['content-type'] = 'application/json; charset=utf-8'
+        self.assertTrue(is_json_payload(http_entity))
+
+        headers['content-type'] = 'application/xml; charset=utf-8'
+        self.assertFalse(is_json_payload(http_entity))

src/azure_devtools/scenario_tests/utilities.py

@@ -34,3 +34,32 @@ def get_sha1_hash(file_path):
             sha1.update(data)
 
     return sha1.hexdigest()
+
+
+def _get_content_type(entity):
+    # 'headers' is a field of 'request', but it is a dict-key in 'response'
+    headers = getattr(entity, 'headers', None)
+    if headers is None:
+        headers = entity.get('headers')
+
+    content_type = None
+    if headers:
+        content_type = headers.get('content-type', None)
+        if content_type:
+            # content-type could an array from response, let us extract it out
+            content_type = content_type[0] if isinstance(content_type, list) else content_type
+            content_type = content_type.split(";")[0].lower()
+    return content_type
+
+
+def is_text_payload(entity):
+    text_content_list = ['application/json', 'application/xml', 'text/', 'application/test-content']
+
+    content_type = _get_content_type(entity)
+    if content_type:
+        return any(content_type.startswith(x) for x in text_content_list)
+    return True
+
+
+def is_json_payload(entity):
+    return _get_content_type(entity) == 'application/json'