Merge pull request #82 from garnaat/use-placebo

Use placebo

requirements.txt

@@ -1,4 +1,5 @@
 boto3>=1.2.3
+placebo==0.4.3
 six>=1.9.0
 PyYAML==3.11
 python-dateutil>=2.1,<3.0.0

skew/__init__.py

@@ -18,7 +18,7 @@
 __version__ = open(os.path.join(os.path.dirname(__file__), '_version')).read()
 
 
-def scan(sku, aws_creds=None):
+def scan(sku, **kwargs):
     """
     Scan (i.e. look up) a SKU.
 
@@ -37,4 +37,4 @@ def scan(sku, aws_creds=None):
     but since there is currently only one (ARN) let's not over-complicate
     things.
     """
-    return ARN(sku, aws_creds=aws_creds)
+    return ARN(sku, **kwargs)

skew/arn/__init__.py

@@ -108,7 +108,7 @@ def choices(self, context=None):
             all_resources = ['*']
         return all_resources
 
-    def enumerate(self, context, aws_creds=None):
+    def enumerate(self, context, **kwargs):
         LOG.debug('Resource.enumerate %s', context)
         _, provider, service_name, region, account = context
         resource_type, resource_id = self._split_resource(self.pattern)
@@ -119,7 +119,7 @@ def enumerate(self, context, aws_creds=None):
             resource_path = '.'.join([provider, service_name, resource_type])
             resource_cls = skew.resources.find_resource_class(resource_path)
             resources.extend(resource_cls.enumerate(
-                self._arn, region, account, resource_id, aws_creds=aws_creds))
+                self._arn, region, account, resource_id, **kwargs))
         return resources
 
 
@@ -132,12 +132,12 @@ def __init__(self, pattern, arn):
     def choices(self, context=None):
         return list(self._accounts.keys())
 
-    def enumerate(self, context, aws_creds=None):
+    def enumerate(self, context, **kwargs):
         LOG.debug('Account.enumerate %s', context)
         for match in self.matches(context):
             context.append(match)
             for resource in self._arn.resource.enumerate(
-                    context, aws_creds=aws_creds):
+                    context, **kwargs):
                 yield resource
             context.pop()
 
@@ -179,12 +179,12 @@ def choices(self, context=None):
         return self._service_region_map.get(
             service, self._all_region_names)
 
-    def enumerate(self, context, aws_creds=None):
+    def enumerate(self, context, **kwargs):
         LOG.debug('Region.enumerate %s', context)
         for match in self.matches(context):
             context.append(match)
             for account in self._arn.account.enumerate(
-                    context, aws_creds=aws_creds):
+                    context, **kwargs):
                 yield account
             context.pop()
 
@@ -198,12 +198,12 @@ def choices(self, context=None):
             provider = self._arn.provider.pattern
         return skew.resources.all_services(provider)
 
-    def enumerate(self, context, aws_creds=None):
+    def enumerate(self, context, **kwargs):
         LOG.debug('Service.enumerate %s', context)
         for match in self.matches(context):
             context.append(match)
             for region in self._arn.region.enumerate(
-                    context, aws_creds=aws_creds):
+                    context, **kwargs):
                 yield region
             context.pop()
 
@@ -213,12 +213,12 @@ class Provider(ARNComponent):
     def choices(self, context=None):
         return ['aws']
 
-    def enumerate(self, context, aws_creds=None):
+    def enumerate(self, context, **kwargs):
         LOG.debug('Provider.enumerate %s', context)
         for match in self.matches(context):
             context.append(match)
             for service in self._arn.service.enumerate(
-                    context, aws_creds=aws_creds):
+                    context, **kwargs):
                 yield service
             context.pop()
 
@@ -228,12 +228,12 @@ class Scheme(ARNComponent):
     def choices(self, context=None):
         return ['arn']
 
-    def enumerate(self, context, aws_creds=None):
+    def enumerate(self, context, **kwargs):
         LOG.debug('Scheme.enumerate %s', context)
         for match in self.matches(context):
             context.append(match)
             for provider in self._arn.provider.enumerate(
-                    context, aws_creds=aws_creds):
+                    context, **kwargs):
                 yield provider
             context.pop()
 
@@ -242,11 +242,11 @@ class ARN(object):
 
     ComponentClasses = [Scheme, Provider, Service, Region, Account, Resource]
 
-    def __init__(self, arn_string='arn:aws:*:*:*:*', aws_creds=None):
+    def __init__(self, arn_string='arn:aws:*:*:*:*', **kwargs):
         self.query = None
         self._components = None
         self._build_components_from_string(arn_string)
-        self.aws_creds = aws_creds
+        self.kwargs = kwargs
 
     def __repr__(self):
         return ':'.join([str(c) for c in self._components])
@@ -308,5 +308,5 @@ def resource(self):
 
     def __iter__(self):
         context = []
-        for scheme in self.scheme.enumerate(context, aws_creds=self.aws_creds):
+        for scheme in self.scheme.enumerate(context, **self.kwargs):
             yield scheme

skew/awsclient.py

@@ -13,12 +13,12 @@
 # limitations under the License.
 
 import logging
-import json
-import os
+import time
 
 import datetime
 import jmespath
 import boto3
+from botocore.exceptions import ClientError
 
 from skew.config import get_config
 
@@ -35,18 +35,20 @@ def json_encoder(obj):
 
 class AWSClient(object):
 
-    def __init__(self, service_name, region_name, account_id, aws_creds=None):
+    def __init__(self, service_name, region_name, account_id, **kwargs):
         self._config = get_config()
         self._service_name = service_name
         self._region_name = region_name
         self._account_id = account_id
         self._has_credentials = False
-        if not aws_creds:
-            # If no creds, need profile name to retrieve creds from ~/.aws/credentials
+        self.aws_creds = kwargs.get('aws_creds')
+        if self.aws_creds is None:
+            # no aws_creds, need profile to get creds from ~/.aws/credentials
             self._profile = self._config['accounts'][account_id]['profile']
-        self.aws_creds = aws_creds
+        self.placebo = kwargs.get('placebo')
+        self.placebo_dir = kwargs.get('placebo_dir')
+        self.placebo_mode = kwargs.get('placebo_mode', 'record')
         self._client = self._create_client()
-        self._record_path = self._config.get('record_path', None)
 
     @property
     def service_name(self):
@@ -64,49 +66,19 @@ def account_id(self):
     def profile(self):
         return self._profile
 
-    def _record(self, op_name, kwargs, data):
-        """
-        This is a little hack to enable easier unit testing of the code.
-        Since botocore has its own set of tests, I'm not interested in
-        trying to test it again here.  So, this recording capability allows
-        us to save the data coming back from botocore as JSON files which
-        can then be used by the mocked awsclient in the unit test directory.
-        To enable this, add something like this to your skew config file:
-
-              record_path: ~/projects/skew/skew/tests/unit/data
-
-        and the JSON data files will get stored in this path.
-        """
-        if self._record_path:
-            path = os.path.expanduser(self._record_path)
-            path = os.path.expandvars(path)
-            path = os.path.join(path, self.service_name)
-            if not os.path.isdir(path):
-                os.mkdir(path)
-            path = os.path.join(path, self.region_name)
-            if not os.path.isdir(path):
-                os.mkdir(path)
-            path = os.path.join(path, self.account_id)
-            if not os.path.isdir(path):
-                os.mkdir(path)
-            filename = op_name
-            if kwargs:
-                for k, v in kwargs.items():
-                    if k != 'query':
-                        filename += '_{}_{}'.format(k, v)
-            filename += '.json'
-            path = os.path.join(path, filename)
-            with open(path, 'wb') as fp:
-                json.dump(data, fp, indent=4, default=json_encoder,
-                          ensure_ascii=False)
-
     def _create_client(self):
         if self.aws_creds:
             session = boto3.Session(**self.aws_creds)
         else:
             session = boto3.Session(
-                profile_name=self.profile, region_name=self.region_name)
-        return session.client(self.service_name)
+                profile_name=self.profile)
+        if self.placebo and self.placebo_dir:
+            pill = self.placebo.attach(session, self.placebo_dir)
+            if self.placebo_mode == 'record':
+                pill.record()
+            elif self.placebo_mode == 'playback':
+                pill.playback()
+        return session.client(self.service_name, region_name=self.region_name)
 
     def call(self, op_name, query=None, **kwargs):
         """
@@ -144,12 +116,21 @@ def call(self, op_name, query=None, **kwargs):
             data = results.build_full_result()
         else:
             op = getattr(self._client, op_name)
-            data = op(**kwargs)
+            done = False
+            data = {}
+            while not done:
+                try:
+                    data = op(**kwargs)
+                    done = True
+                except ClientError as e:
+                    if 'Throttling' in str(e):
+                        time.sleep(1)
+                except Exception:
+                    done = True
         if query:
             data = query.search(data)
-        self._record(op_name, kwargs, data)
         return data
 
 
-def get_awsclient(service_name, region_name, account_id, aws_creds=None):
-    return AWSClient(service_name, region_name, account_id, aws_creds)
+def get_awsclient(service_name, region_name, account_id, **kwargs):
+    return AWSClient(service_name, region_name, account_id, **kwargs)

skew/resources/aws/cloudformation.py

@@ -18,9 +18,9 @@
 class Stack(AWSResource):
 
     @classmethod
-    def enumerate(cls, arn, region, account, resource_id=None, aws_creds=None):
+    def enumerate(cls, arn, region, account, resource_id=None, **kwargs):
         resources = super(Stack, cls).enumerate(arn, region, account,
-                                                resource_id, aws_creds)
+                                                resource_id, **kwargs)
         for stack in resources:
             stack.data['Resources'] = []
             for stack_resource in stack:

skew/resources/aws/lambda.py

@@ -23,9 +23,9 @@
 class Function(AWSResource):
 
     @classmethod
-    def enumerate(cls, arn, region, account, resource_id=None):
+    def enumerate(cls, arn, region, account, resource_id=None, **kwargs):
         resources = super(Function, cls).enumerate(arn, region, account,
-                                                   resource_id)
+                                                   resource_id, **kwargs)
         for r in resources:
             r.data['EventSources'] = []
             kwargs = {'FunctionName': r.data['FunctionName']}

skew/resources/aws/s3.py

@@ -23,10 +23,10 @@ class Bucket(AWSResource):
     _location_cache = {}
 
     @classmethod
-    def enumerate(cls, arn, region, account, resource_id=None, aws_creds=None):
+    def enumerate(cls, arn, region, account, resource_id=None, **kwargs):
         resources = super(Bucket, cls).enumerate(arn, region, account,
                                                  resource_id,
-                                                 aws_creds=aws_creds)
+                                                 **kwargs)
         region_resources = []
         if region is None:
             region = 'us-east-1'

skew/resources/resource.py

@@ -23,9 +23,9 @@
 class Resource(object):
 
     @classmethod
-    def enumerate(cls, arn, region, account, resource_id=None, aws_creds=None):
+    def enumerate(cls, arn, region, account, resource_id=None, **kwargs):
         client = skew.awsclient.get_awsclient(
-            cls.Meta.service, region, account, aws_creds=aws_creds)
+            cls.Meta.service, region, account, **kwargs)
         kwargs = {}
         do_client_side_filtering = False
         if resource_id and resource_id != '*':