address review feedback

Author: yugangw-msft

src/azure_devtools/scenario_tests/__init__.py

@@ -5,9 +5,9 @@
 
 from .base import IntegrationTestBase, ReplayableTest, LiveTest
 from .exceptions import AzureTestError
-from .decorators import live_only, record_only
+from .decorators import live_only, record_only, AllowLargeResponse
 from .patches import mock_in_unit_test, patch_time_sleep_api, patch_long_run_operation_delay
-from .preparers import AbstractPreparer, SingleValueReplacer, AllowLargeResponse
+from .preparers import AbstractPreparer, SingleValueReplacer
 from .recording_processors import (
     RecordingProcessor, SubscriptionRecordingProcessor,
     LargeRequestBodyProcessor, LargeResponseBodyProcessor, LargeResponseBodyReplacer,

src/azure_devtools/scenario_tests/decorators.py

@@ -4,8 +4,10 @@
 # --------------------------------------------------------------------------------------------
 
 import os
+import functools
 import unittest
 from .const import ENV_LIVE_TEST
+from .utilities import trim_kwargs_from_test_function
 
 
 def live_only():
@@ -18,3 +20,25 @@ def record_only():
     return unittest.skipUnless(
         not os.environ.get(ENV_LIVE_TEST, False),
         'This test is excluded from being run live. To force a recording, please remove the recording file.')
+
+
+class AllowLargeResponse(object):  # pylint: disable=too-few-public-methods
+
+    def __init__(self, size_kb=1024):
+        self.size_kb = size_kb
+
+    def __call__(self, fn):
+        def _preparer_wrapper(test_class_instance, **kwargs):
+            from azure_devtools.scenario_tests import LargeResponseBodyProcessor
+            large_resp_body = next((r for r in 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
+
+            trim_kwargs_from_test_function(fn, kwargs)
+
+            fn(test_class_instance, **kwargs)
+
+        setattr(_preparer_wrapper, '__is_preparer', True)
+        functools.update_wrapper(_preparer_wrapper, fn)
+        return _preparer_wrapper

src/azure_devtools/scenario_tests/preparers.py

@@ -4,11 +4,10 @@
 # --------------------------------------------------------------------------------------------
 
 import contextlib
-import inspect
 import functools
 
 from .base import ReplayableTest
-from .utilities import create_random_name, is_text_payload
+from .utilities import create_random_name, is_text_payload, trim_kwargs_from_test_function
 from .recording_processors import RecordingProcessor
 
 
@@ -48,7 +47,7 @@ def _preparer_wrapper(test_class_instance, **kwargs):
             if parameter_update:
                 kwargs.update(parameter_update)
 
-            _trim_kwargs_from_test_function(fn, kwargs)
+            trim_kwargs_from_test_function(fn, kwargs)
 
             fn(test_class_instance, **kwargs)
 
@@ -120,41 +119,3 @@ def process_response(self, response):
         self.replace_header(response, 'azure-asyncoperation', self.random_name, self.moniker)
 
         return response
-
-
-class AllowLargeResponse(object):  # pylint: disable=too-few-public-methods
-
-    def __init__(self, size_kb=1024):
-        self.size_kb = size_kb
-
-    def __call__(self, fn):
-        def _preparer_wrapper(test_class_instance, **kwargs):
-            from azure_devtools.scenario_tests import LargeResponseBodyProcessor
-            large_resp_body = next((r for r in 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
-
-            _trim_kwargs_from_test_function(fn, kwargs)
-
-            fn(test_class_instance, **kwargs)
-
-        setattr(_preparer_wrapper, '__is_preparer', True)
-        functools.update_wrapper(_preparer_wrapper, fn)
-        return _preparer_wrapper
-
-# Utility
-
-def _trim_kwargs_from_test_function(fn, kwargs):
-    # the next function is the actual test function. the kwargs need to be trimmed so
-    # that parameters which are not required will not be passed to it.
-    if not is_preparer_func(fn):
-        args, _, kw, _ = inspect.getargspec(fn)  # pylint: disable=deprecated-method
-        if kw is None:
-            args = set(args)
-            for key in [k for k in kwargs if k not in args]:
-                del kwargs[key]
-
-
-def is_preparer_func(fn):
-    return getattr(fn, '__is_preparer', False)

src/azure_devtools/scenario_tests/recording_processors.py

@@ -87,7 +87,7 @@ def process_response(self, response):
             if length > self._max_response_body * 1024:
 
                 if is_json_payload(response):
-                    from .preparers import AllowLargeResponse  # pylint: disable=cyclic-import
+                    from .decorators 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__))

src/azure_devtools/scenario_tests/utilities.py

@@ -7,6 +7,7 @@
 import math
 import os
 import base64
+import inspect
 
 
 def create_random_name(prefix='aztest', length=24):
@@ -63,3 +64,18 @@ def is_text_payload(entity):
 
 def is_json_payload(entity):
     return _get_content_type(entity) == 'application/json'
+
+
+def trim_kwargs_from_test_function(fn, kwargs):
+    # the next function is the actual test function. the kwargs need to be trimmed so
+    # that parameters which are not required will not be passed to it.
+    if not is_preparer_func(fn):
+        args, _, kw, _ = inspect.getargspec(fn)  # pylint: disable=deprecated-method
+        if kw is None:
+            args = set(args)
+            for key in [k for k in kwargs if k not in args]:
+                del kwargs[key]
+
+
+def is_preparer_func(fn):
+    return getattr(fn, '__is_preparer', False)