Fix stats error

Author: whu-lyft

kmsauth/__init__.py

@@ -169,13 +169,17 @@ def _format_auth_key(self, keys):
         )
 
     def _get_key_arn(self, key):
-        logging.debug(f'Getting key ARN for {key}')
         if key.startswith('arn:aws:kms:'):
             self.KEY_METADATA[key] = {
                 'KeyMetadata': {'Arn': key}
             }
         if key not in self.KEY_METADATA:
-            with self.stats.timer('kms_describe_key'):
+            if self.stats:
+                with self.stats.timer('kms_describe_key'):
+                    self.KEY_METADATA[key] = self.kms_client.describe_key(
+                        KeyId='{0}'.format(key)
+                    )
+            else:
                 self.KEY_METADATA[key] = self.kms_client.describe_key(
                     KeyId='{0}'.format(key)
                 )
@@ -261,10 +265,11 @@ def decrypt_token(self, username, token):
         except Exception:
             raise TokenValidationError('Authentication error.')
         if token_key not in self.TOKENS:
-            self.stats.incr('token_cache.miss')
-            self.stats.gauge('token_cache.size_at_miss', len(self.TOKENS))
-            if len(self.TOKENS) >= self.token_cache_size:
-                self.stats.incr('token_cache.eviction')
+            if self.stats:
+                self.stats.incr('token_cache.miss')
+                self.stats.gauge('token_cache.size_at_miss', len(self.TOKENS))
+                if len(self.TOKENS) >= token_cache_size:
+                    self.stats.incr('token_cache.eviction')
 
             try:
                 token = base64.b64decode(token)
@@ -323,7 +328,8 @@ def decrypt_token(self, username, token):
                     'Authentication error. General error.'
                 )
         else:
-            self.stats.incr('token_cache.hit')
+            if self.stats:
+                self.stats.incr('token_cache.hit')
             ret = self.TOKENS[token_key]
         now = datetime.datetime.utcnow()
         try: