Clean up; add test for parsing AWS response

Author: swfunc

aws/logs_monitoring/enhanced_metrics.py

@@ -1,14 +1,12 @@
-import re
-import os
 import logging
+import os
+import re
 
 from collections import defaultdict
 from time import time
 
-
 import boto3
 from botocore.exceptions import ClientError
-
 from datadog_lambda.metric import lambda_metric
 
 
@@ -32,7 +30,8 @@
 log = logging.getLogger()
 log.setLevel(logging.INFO)
 
-def set_log_level_to_env_var:
+
+def set_log_level_to_env_var():
  """Reads the log level env var and sets the log level according to its value
 
  Defaults to INFO level
@@ -45,6 +44,7 @@ def set_log_level_to_env_var:
  if env_var_log_level == "warn" or env_var_log_level == "warning":
  log.setLevel(logging.WARN)
 
+
 set_log_level_to_env_var()
 
 
@@ -86,7 +86,15 @@ def get_dd_tag_string_from_aws_dict(aws_key_value_tag_dict):
 def parse_get_resources_response_for_tags_by_arn(get_resources_page):
  """Parses a page of GetResources response for the mapping from ARN to tags
 
+ Args:
+ get_resources_page (dict<str, multiple types>): one page of the GetResources response. Ex:
+ [{
+ 'ResourceARN': 'arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda',
+ 'Tags': [{'Key': 'stage', 'Value': 'dev'}, {'Key': 'team', 'Value': 'serverless'}]
+ }]
 
+ Returns:
+ tags_by_arn (dict<str, str[]>): Lambda tag lists keyed by ARN
  """
  tags_by_arn = defaultdict(list)
 
@@ -180,6 +188,7 @@ def get_lambda_tags(self, resource_arn):
 
  return function_tags
 
+
 # Store the cache in the global scope so that it will be reused as long as
 # the log forwarder Lambda container is running
 account_lambda_tags_cache = LambdaTagsCache()
@@ -236,11 +245,15 @@ def parse_and_submit_enhanced_metrics(logs):
  """
  # Wrap everything in try/catch to prevent failing the Lambda on enhanced metrics
  try:
- enhanced_metrics = generate_enhanced_lambda_metrics(logs, account_lambda_tags_cache)
+ enhanced_metrics = generate_enhanced_lambda_metrics(
+ logs, account_lambda_tags_cache
+ )
  for enhanced_metric in enhanced_metrics:
  enhanced_metric.submit_to_dd()
  except Exception:
- log.exception("Encountered an error while trying to parse and submit enhanced metrics")
+ log.exception(
+ "Encountered an error while trying to parse and submit enhanced metrics"
+ )
 
 
 def generate_enhanced_lambda_metrics(logs, tags_cache):
@@ -283,12 +296,7 @@ def generate_enhanced_lambda_metrics(logs, tags_cache):
 
  # If the log dict is missing any of this data it's not a Lambda REPORT log and we move on
  if not all(
- (
- log_function_arn,
- log_message,
- timestamp,
- log_message.startswith("REPORT"),
- )
+ (log_function_arn, log_message, timestamp, log_message.startswith("REPORT"))
  ):
  continue
 
@@ -313,6 +321,7 @@ def generate_enhanced_lambda_metrics(logs, tags_cache):
 
  return enhanced_metrics
 
+
 # Names to use for metrics and for the named regex groups
 REQUEST_ID_FIELD_NAME = "request_id"
 DURATION_METRIC_NAME = "duration"
@@ -364,9 +373,9 @@ def parse_lambda_tags_from_arn(arn):
  _, _, _, region, account_id, _, function_name = split_arn
 
  return [
- "functionname:{}".format(function_name),
- "account_id:{}".format(account_id),
  "region:{}".format(region),
+ "account_id:{}".format(account_id),
+ "functionname:{}".format(function_name),
  ]
 
 

aws/logs_monitoring/lambda_function.py

@@ -16,7 +16,7 @@
 import ssl
 import six.moves.urllib as urllib # for for Python 2.7 urllib.unquote_plus
 import itertools
-import uuid
+
 from io import BytesIO, BufferedReader
 
 from enhanced_metrics import parse_and_submit_enhanced_metrics
@@ -59,11 +59,6 @@
  DD_URL = os.getenv("DD_URL", default="lambda-http-intake.logs." + DD_SITE)
 
 
-lambda_container_uuid = uuid.uuid4()
-container_start_time = time.time()
-
-
-
 class ScrubbingRuleConfig(object):
  def __init__(self, name, pattern, placeholder):
  self.name = name

aws/logs_monitoring/tests/test_enhanced_metrics.py

@@ -8,10 +8,10 @@
  parse_lambda_tags_from_arn,
  generate_enhanced_lambda_metrics,
  LambdaTagsCache,
+ parse_get_resources_response_for_tags_by_arn,
 )
 
 
-# python -m unittest tests.test_enhanced_metrics from parent dir
 class TestEnhancedMetrics(unittest.TestCase):
 
  maxDiff = None
@@ -47,6 +47,42 @@ def test_parse_lambda_tags_from_arn(self):
  ],
  )
 
+ def test_generate_custom_tags_cache(self):
+ self.assertEqual(
+ parse_get_resources_response_for_tags_by_arn(
+ {
+ "ResourceTagMappingList": [
+ {
+ "ResourceARN": "arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda-dev",
+ "Tags": [
+ {"Key": "stage", "Value": "dev"},
+ {"Key": "team", "Value": "serverless"},
+ ],
+ },
+ {
+ "ResourceARN": "arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda-prod",
+ "Tags": [
+ {"Key": "stage", "Value": "prod"},
+ {"Key": "team", "Value": "serverless"},
+ {"Key": "datacenter", "Value": "eu"},
+ ],
+ },
+ ]
+ }
+ ),
+ {
+ "arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda-dev": [
+ "stage:dev",
+ "team:serverless",
+ ],
+ "arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda-prod": [
+ "stage:prod",
+ "team:serverless",
+ "datacenter:eu",
+ ],
+ },
+ )
+
  def test_parse_metrics_from_report_log(self):
  parsed_metrics = parse_metrics_from_report_log(self.malformed_report)
  self.assertEqual(parsed_metrics, [])
@@ -115,7 +151,7 @@ def test_parse_metrics_from_report_log(self):
  ],
  )
 
- @patch("enhanced_metrics.build_arn_to_lambda_tags_cache")
+ @patch("enhanced_metrics.build_tags_by_arn_cache")
  def test_generate_enhanced_lambda_metrics(self, mock_build_cache):
  mock_build_cache.return_value = {}
  tags_cache = LambdaTagsCache()