Record based on parameter hashing

placebo/pill.py

@@ -11,12 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import hashlib
+import json
 import os
-import glob
 import re
+import copy
 import uuid
 import logging
+from typing import List, Dict
 
 from placebo.serializer import Format, get_deserializer, get_serializer
 
@@ -30,8 +32,14 @@ def __init__(self, status_code):
         self.status_code = status_code
 
 
-class Pill(object):
+def _filter_hash(unfiltered: Dict, keys: List[str]) -> Dict:
+    filtered = {}
+    for key in keys:
+        filtered[key] = unfiltered[key]
+    return filtered
 
+
+class Pill(object):
     clients = []
 
     def __init__(self, prefix=None, debug=False, record_format=Format.JSON):
@@ -44,16 +52,17 @@ def __init__(self, prefix=None, debug=False, record_format=Format.JSON):
             record_format = Format.DEFAULT
 
         self._serializer = get_serializer(record_format)
-        self._filename_re = re.compile(r'.*\..*_(?P<index>\d+).{0}'.format(
+        self._filename_re = re.compile(r'.*\..*.{0}'.format(
             record_format))
         self._record_format = record_format
         self.prefix = prefix
         self._uuid = str(uuid.uuid4())
         self._data_path = None
         self._mode = None
+        self._services = ""
+        self._operations = ""
         self._session = None
-        self._index = {}
-        self.events = []
+        self.events = {}
         self.clients = []
 
     @property
@@ -142,82 +151,96 @@ def attach(self, session, data_path):
         self._data_path = data_path
         session.events.register('creating-client-class', self._create_client)
 
-    def record(self, services='*', operations='*'):
+    def _registered(self, service_name: str, operation_name: str) -> bool:
+        """
+        Returns true if the service_name and operation_name match the regex
+        pattern used with record(). This is meant to be called from the handlers
+        as a workaround for boto3 not supporting globbing natively.
+        """
+        if re.match(self._services, service_name) and re.match(self._operations, operation_name):
+            LOG.debug('Handling Event: %s -- %s', service_name, operation_name)
+            return True
+        else:
+            LOG.debug('Ignoring Event: %s -- %s', service_name, operation_name)
+            return False
+
+    def record(self, services='.*', operations='.*'):
+        """
+        Record start's recording event's to the filesystem. You can use regex
+        expressions for services and operations, however the filtering is done
+        in the handler's itself.
+        """
         if self._mode == 'playback':
             self.stop()
         self._mode = 'record'
-        for service in services.split(','):
-            for operation in operations.split(','):
-                event = 'after-call.{0}.{1}'.format(
-                    service.strip(), operation.strip())
-                LOG.debug('recording: %s', event)
-                self.events.append(event)
-                self._session.events.register(
-                    event, self._record_data, 'placebo-record-mode')
-                for client in self.clients:
-                    client.meta.events.register(
-                        event, self._record_data, 'placebo-record-mode')
+
+        # You can't register event's with normal globbing, but we can store our original
+        # pattern and check it in the handler to get the same behavior.
+        self._services = services
+        self._operations = operations
+
+        self._register('record', 'before-call', "*", "*", self._record_params)
+        self._register('record', 'after-call', "*", "*", self._record_data)
+
+    def _register(self, mode, event_name, operation, service, function):
+        name = '{0}.{1}.{2}'.format(event_name, service.strip(), operation.strip())
+        unique_id = 'placebo-{}-mode-{}-{}'.format(mode, event_name, function.__name__)
+        LOG.debug('recording: %s -- %s', name, unique_id)
+        self.events[name] = unique_id
+        self._session.events.register(name, function, unique_id)
+        for client in self.clients:
+            client.meta.events.register(name, function, unique_id)
 
     def playback(self):
         if self.mode == 'record':
             self.stop()
         if self.mode is None:
-            event = 'before-call.*.*'
-            self.events.append(event)
-            self._session.events.register(
-                event, self._mock_request, 'placebo-playback-mode')
             self._mode = 'playback'
-            for client in self.clients:
-                client.meta.events.register(
-                    event, self._mock_request, 'placebo-playback-mode')
+            self._register('playback', 'before-call', '*', '*', self._record_params)
+            self._register('playback', 'before-call', '*', '*', self._mock_request)
 
     def stop(self):
         LOG.debug('stopping, mode=%s', self.mode)
-        if self.mode == 'record':
-            if self._session:
-                for event in self.events:
-                    self._session.events.unregister(
-                        event, unique_id='placebo-record-mode')
-                    for client in self.clients:
-                        client.meta.events.unregister(
-                            event, unique_id='placebo-record-mode')
-                self.events = []
-        elif self.mode == 'playback':
-            if self._session:
-                for event in self.events:
-                    self._session.events.unregister(
-                        event, unique_id='placebo-playback-mode')
-                    for client in self.clients:
-                        client.meta.events.unregister(
-                            event, unique_id='placebo-playback-mode')
-                self.events = []
+        if self._session:
+            for name, unique_id in self.events.items():
+                self._session.events.unregister(name, unique_id=unique_id)
+                for client in self.clients:
+                    client.meta.events.unregister(name, unique_id=unique_id)
+            self.events = {}
         self._mode = None
 
-    def _record_data(self, http_response, parsed, model, **kwargs):
-        LOG.debug('_record_data')
+    def _record_params(self, params, context, model, **kwargs):
         service_name = model.service_model.endpoint_prefix
         operation_name = model.name
-        self.save_response(service_name, operation_name, parsed,
-                           http_response.status_code)
+        if not self._registered(service_name, operation_name):
+            return
+
+        p = copy.deepcopy(params)  # Be careful not to modify the original
+        p = _filter_hash(p, ["url_path", "query_string", "method", "body", "url"])
+        p["client_region"] = context["client_region"]
+        params_h = json.dumps(p).encode()
+        context["_pill_params_hash"] = hashlib.sha256(params_h).hexdigest()
+        LOG.debug('_record_params')
 
-    def get_new_file_path(self, service, operation):
-        base_name = '{0}.{1}'.format(service, operation)
+    def _record_data(self, http_response, parsed, model, context, **kwargs):
+        LOG.debug('_record_data')
+        service_name = model.service_model.endpoint_prefix
+        operation_name = model.name
+        if not self._registered(service_name, operation_name):
+            return
+        params_hash = context["_pill_params_hash"]
+        self.save_response(service_name, operation_name, params_hash,
+                           parsed, http_response.status_code)
+
+    def get_new_file_path(self, service, operation, params_hash):
+        base_name = '{0}.{1}.{2}'.format(service, operation, params_hash)
         if self.prefix:
             base_name = '{0}.{1}'.format(self.prefix, base_name)
-        LOG.debug('get_new_file_path: %s', base_name)
-        index = 0
-        glob_pattern = os.path.join(self._data_path, base_name + '_*')
-        for file_path in glob.glob(glob_pattern):
-            file_name = os.path.basename(file_path)
-            m = self._filename_re.match(file_name)
-            if m:
-                i = int(m.group('index'))
-                if i > index:
-                    index = i
-        index += 1
+
         return os.path.join(
-            self._data_path, '{0}_{1}.{2}'.format(
-                base_name, index, self.record_format))
+            self._data_path,
+            '{0}.{1}'.format(base_name, self.record_format)
+        )
 
     @staticmethod
     def find_file_format(file_name):
@@ -230,36 +253,26 @@ def find_file_format(file_name):
                 return file_path, file_format
         return None, None
 
-    def get_next_file_path(self, service, operation):
+    def get_next_file_path(self, service, operation, params_hash):
         """
         Returns a tuple with the next file to read and the serializer
         format used
         """
-        base_name = '{0}.{1}'.format(service, operation)
+        base_name = '{0}.{1}.{2}'.format(service, operation, params_hash)
         if self.prefix:
             base_name = '{0}.{1}'.format(self.prefix, base_name)
-        LOG.debug('get_next_file_path: %s', base_name)
-        next_file = None
-        serializer_format = None
-        index = self._index.setdefault(base_name, 1)
-
-        while not next_file:
-            file_name = os.path.join(
-                self._data_path, base_name + '_{0}'.format(index))
-            next_file, serializer_format = self.find_file_format(file_name)
-            if next_file:
-                self._index[base_name] += 1
-            elif index != 1:
-                index = 1
-                self._index[base_name] = 1
-            else:
-                raise IOError('response file ({0}.[{1}]) not found'.format(
-                    file_name, "|".join(Format.ALLOWED)))
+        file_path = os.path.join(self.data_path, base_name)
+        LOG.debug('get_next_file_path: %s', file_path)
+
+        next_file, serializer_format = self.find_file_format(file_path)
+        if not next_file:
+            raise IOError('response file ({0}.[{1}]) not found'.format(
+                file_path, "|".join(Format.ALLOWED)))
 
         return next_file, serializer_format
 
-    def save_response(self, service, operation, response_data,
-                      http_response=200):
+    def save_response(self, service, operation, params_hash,
+                      response_data, http_response=200):
         """
         Store a response to the data directory.  The ``operation``
         should be the name of the operation in the service API (e.g.
@@ -269,31 +282,40 @@ def save_response(self, service, operation, response_data,
         multiple responses for a given operation and they will be
         returned in order.
         """
-        LOG.debug('save_response: %s.%s', service, operation)
-        filepath = self.get_new_file_path(service, operation)
+        LOG.debug('save_response: %s.%s.%s', service, operation, params_hash)
+
+        # TODO: tests
+        filepath = self.get_new_file_path(service, operation, params_hash)
+
         LOG.debug('save_response: path=%s', filepath)
         data = {'status_code': http_response,
                 'data': response_data}
+
+        dir = os.path.dirname(filepath)
+        if not os.path.exists(dir):
+            os.mkdir(dir)
+
         with open(filepath, Format.write_mode(self.record_format)) as fp:
             self._serializer(data, fp)
 
-    def load_response(self, service, operation):
-        LOG.debug('load_response: %s.%s', service, operation)
+    def load_response(self, service, operation, params_hash):
+        LOG.debug('load_response: %s.%s.%s', service, operation, params_hash)
         (response_file, file_format) = self.get_next_file_path(
-            service, operation)
+            service, operation, params_hash)
         LOG.debug('load_responses: %s', response_file)
         with open(response_file, Format.read_mode(file_format)) as fp:
             response_data = get_deserializer(file_format)(fp)
         return (FakeHttpResponse(response_data['status_code']),
                 response_data['data'])
 
-    def _mock_request(self, **kwargs):
+    def _mock_request(self, context, params, **kwargs):
         """
         A mocked out make_request call that bypasses all network calls
         and simply returns any mocked responses defined.
         """
         model = kwargs.get('model')
         service = model.service_model.endpoint_prefix
         operation = model.name
-        LOG.debug('_make_request: %s.%s', service, operation)
-        return self.load_response(service, operation)
+        params_hash = context["_pill_params_hash"]
+        LOG.debug('_make_request: %s.%s.%s', service, operation, params_hash)
+        return self.load_response(service, operation, params_hash)

requirements.txt

@@ -1,2 +1,5 @@
 boto3>=1.8.6
 nose>=1.3.7
+
+mock~=4.0.3
+moto==1.3.14

tests/unit/responses/saved/iam.GetPolicy.7b1fd414173035b6f7f70aee6c4aedafd78910bd8632d9f0211057e5fd90bef7_1.json

@@ -0,0 +1,47 @@
+{
+    "status_code": 200,
+    "data": {
+        "Policy": {
+            "PolicyName": "AWSLambdaFullAccess",
+            "PolicyId": "ANPAI6E2CYYMI4XI7AA5K",
+            "Arn": "arn:aws:iam::aws:policy/AWSLambdaFullAccess",
+            "Path": "/",
+            "DefaultVersionId": "v8",
+            "AttachmentCount": 0,
+            "PermissionsBoundaryUsageCount": 0,
+            "IsAttachable": true,
+            "Description": "Provides full access to Lambda, S3, DynamoDB, CloudWatch Metrics and Logs.",
+            "CreateDate": {
+                "__class__": "datetime",
+                "year": 2015,
+                "month": 2,
+                "day": 6,
+                "hour": 18,
+                "minute": 40,
+                "second": 45,
+                "microsecond": 0
+            },
+            "UpdateDate": {
+                "__class__": "datetime",
+                "year": 2017,
+                "month": 11,
+                "day": 27,
+                "hour": 23,
+                "minute": 22,
+                "second": 38,
+                "microsecond": 0
+            }
+        },
+        "ResponseMetadata": {
+            "RequestId": "9c5fb25b-001c-46f7-aafb-8dbb89182bcd",
+            "HTTPStatusCode": 200,
+            "HTTPHeaders": {
+                "x-amzn-requestid": "9c5fb25b-001c-46f7-aafb-8dbb89182bcd",
+                "content-type": "text/xml",
+                "content-length": "860",
+                "date": "Sun, 20 Dec 2020 02:45:25 GMT"
+            },
+            "RetryAttempts": 0
+        }
+    }
+}