Use common logging decorator functions

Author: dlzhry2nhs

lambdas/recordprocessor/src/logging_decorator.py

@@ -1,45 +1,16 @@
 """This module contains the logging decorator for sending the appropriate logs to Cloudwatch and Firehose."""
 
-import json
 import os
 import time
 from datetime import datetime
 from functools import wraps
 
-from common.clients import firehose_client, logger
+from common.log_decorator import generate_and_send_logs
 from errors import InvalidHeaders, NoOperationPermissions
 
 STREAM_NAME = os.getenv("SPLUNK_FIREHOSE_NAME", "immunisation-fhir-api-internal-dev-splunk-firehose")
 
 
-def send_log_to_firehose(log_data: dict) -> None:
-    """Sends the log_message to Firehose"""
-    try:
-        record = {"Data": json.dumps({"event": log_data}).encode("utf-8")}
-        response = firehose_client.put_record(DeliveryStreamName=STREAM_NAME, Record=record)
-        logger.info("Log sent to Firehose: %s", response)  # TODO: Should we be logging full response?
-    except Exception as error:  # pylint:disable = broad-exception-caught
-        logger.exception("Error sending log to Firehose: %s", error)
-
-
-def generate_and_send_logs(
-    start_time,
-    base_log_data: dict,
-    additional_log_data: dict,
-    is_error_log: bool = False,
-) -> None:
-    """Generates log data which includes the base_log_data, additional_log_data, and time taken (calculated using the
-    current time and given start_time) and sends them to Cloudwatch and Firehose."""
-    log_data = {
-        **base_log_data,
-        "time_taken": f"{round(time.time() - start_time, 5)}s",
-        **additional_log_data,
-    }
-    log_function = logger.error if is_error_log else logger.info
-    log_function(json.dumps(log_data))
-    send_log_to_firehose(log_data)
-
-
 def file_level_validation_logging_decorator(func):
     """
     Sends the appropriate logs to Cloudwatch and Firehose based on the result of the file_level_validation
@@ -65,7 +36,7 @@ def wrapper(*args, **kwargs):
                 "statusCode": 200,
                 "message": "Successfully sent for record processing",
             }
-            generate_and_send_logs(start_time, base_log_data, additional_log_data=additional_log_data)
+            generate_and_send_logs(STREAM_NAME, start_time, base_log_data, additional_log_data)
             return result
 
         except (InvalidHeaders, NoOperationPermissions, Exception) as e:
@@ -78,7 +49,7 @@ def wrapper(*args, **kwargs):
                 "message": message,
                 "error": str(e),
             }
-            generate_and_send_logs(start_time, base_log_data, additional_log_data, is_error_log=True)
+            generate_and_send_logs(STREAM_NAME, start_time, base_log_data, additional_log_data, is_error_log=True)
             raise
 
     return wrapper

lambdas/recordprocessor/tests/test_logging_decorator.py

@@ -10,12 +10,12 @@
 from boto3 import client as boto3_client
 from moto import mock_firehose, mock_s3
 
-from tests.utils_for_recordprocessor_tests.mock_environment_variables import (
+from utils_for_recordprocessor_tests.mock_environment_variables import (
     MOCK_ENVIRONMENT_DICT,
     BucketNames,
     Firehose,
 )
-from tests.utils_for_recordprocessor_tests.values_for_recordprocessor_tests import (
+from utils_for_recordprocessor_tests.values_for_recordprocessor_tests import (
     MockFileDetails,
     ValidMockFileContent,
 )
@@ -24,10 +24,9 @@
     from common.clients import REGION_NAME
     from errors import InvalidHeaders, NoOperationPermissions
     from file_level_validation import file_level_validation
-    from logging_decorator import generate_and_send_logs, send_log_to_firehose
 
 
-from tests.utils_for_recordprocessor_tests.utils_for_recordprocessor_tests import (
+from utils_for_recordprocessor_tests.utils_for_recordprocessor_tests import (
     GenericSetUp,
     GenericTearDown,
 )
@@ -76,67 +75,6 @@ def run(self, result=None):
                 stack.enter_context(common_patch)
             super().run(result)
 
-    def test_send_log_to_firehose(self):
-        """
-        Tests that the send_log_to_firehose function calls firehose_client.put_record with the correct arguments.
-        NOTE: mock_firehose does not persist the data, so at this level it is only possible to test what the call args
-        were, not that the data reached the destination.
-        """
-        log_data = {"test_key": "test_value"}
-
-        with patch("logging_decorator.firehose_client") as mock_firehose_client:
-            send_log_to_firehose(log_data)
-
-        expected_firehose_record = {"Data": json.dumps({"event": log_data}).encode("utf-8")}
-        mock_firehose_client.put_record.assert_called_once_with(
-            DeliveryStreamName=Firehose.STREAM_NAME, Record=expected_firehose_record
-        )
-
-    def test_generate_and_send_logs(self):
-        """
-        Tests that the generate_and_send_logs function logs the correct data at the correct level for cloudwatch
-        and calls send_log_to_firehose with the correct log data
-        """
-        base_log_data = {"base_key": "base_value"}
-        additional_log_data = {"additional_key": "additional_value"}
-        start_time = 1672531200
-
-        # CASE: Successful log - is_error_log arg set to False
-        with (  # noqa: E999
-            patch("logging_decorator.logger") as mock_logger,  # noqa: E999
-            patch("logging_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
-            patch("logging_decorator.time") as mock_time,  # noqa: E999
-        ):  # noqa: E999
-            mock_time.time.return_value = 1672531200.123456  # Mocks the end time to be 0.123456s after the start time
-            generate_and_send_logs(start_time, base_log_data, additional_log_data, is_error_log=False)
-
-        expected_log_data = {
-            "base_key": "base_value",
-            "time_taken": "0.12346s",
-            "additional_key": "additional_value",
-        }
-        log_data = json.loads(mock_logger.info.call_args[0][0])
-        self.assertEqual(log_data, expected_log_data)
-        mock_send_log_to_firehose.assert_called_once_with(expected_log_data)
-
-        # CASE: Error log - is_error_log arg set to True
-        with (  # noqa: E999
-            patch("logging_decorator.logger") as mock_logger,  # noqa: E999
-            patch("logging_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
-            patch("logging_decorator.time") as mock_time,  # noqa: E999
-        ):  # noqa: E999
-            mock_time.time.return_value = 1672531200.123456  # Mocks the end time to be 0.123456s after the start time
-            generate_and_send_logs(start_time, base_log_data, additional_log_data, is_error_log=True)
-
-        expected_log_data = {
-            "base_key": "base_value",
-            "time_taken": "0.12346s",
-            "additional_key": "additional_value",
-        }
-        log_data = json.loads(mock_logger.error.call_args[0][0])
-        self.assertEqual(log_data, expected_log_data)
-        mock_send_log_to_firehose.assert_called_once_with(expected_log_data)
-
     def test_splunk_logger_successful_validation(self):
         """Tests the splunk logger is called when file-level validation is successful"""
 
@@ -147,12 +85,14 @@ def test_splunk_logger_successful_validation(self):
         )
 
         with (  # noqa: E999
+            patch("common.log_firehose.firehose_client") as mock_firehose_client,  # noqa: E999
+            patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
             patch("logging_decorator.datetime") as mock_datetime,  # noqa: E999
-            patch("logging_decorator.time") as mock_time,  # noqa: E999
-            patch("logging_decorator.logger") as mock_logger,  # noqa: E999
-            patch("logging_decorator.firehose_client") as mock_firehose_client,  # noqa: E999
+            patch("logging_decorator.time") as mock_start_time,  # noqa: E999
+            patch("common.log_decorator.time") as mock_end_time,  # noqa: E999
         ):  # noqa: E999
-            mock_time.time.side_effect = [1672531200, 1672531200.123456]
+            mock_start_time.time.return_value = 1672531200
+            mock_end_time.time.return_value = 1672531200.123456
             mock_datetime.now.return_value = datetime(2024, 1, 1, 12, 0, 0)
             file_level_validation(deepcopy(MOCK_FILE_DETAILS.event_full_permissions_dict))
 
@@ -211,13 +151,15 @@ def test_splunk_logger_handled_failure(self):
                 )
 
                 with (  # noqa: E999
+                    patch("common.log_firehose.firehose_client") as mock_firehose_client,  # noqa: E999
+                    patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
                     patch("logging_decorator.datetime") as mock_datetime,  # noqa: E999
-                    patch("logging_decorator.time") as mock_time,  # noqa: E999
-                    patch("logging_decorator.logger") as mock_logger,  # noqa: E999
-                    patch("logging_decorator.firehose_client") as mock_firehose_client,  # noqa: E999
+                    patch("logging_decorator.time") as mock_start_time,  # noqa: E999
+                    patch("common.log_decorator.time") as mock_end_time,  # noqa: E999
                 ):  # noqa: E999
                     mock_datetime.now.return_value = datetime(2024, 1, 1, 12, 0, 0)
-                    mock_time.time.side_effect = [1672531200, 1672531200.123456]
+                    mock_start_time.time.return_value = 1672531200
+                    mock_end_time.time.return_value = 1672531200.123456
                     with self.assertRaises(expected_error_type):
                         file_level_validation(deepcopy(event_dict))
 
@@ -247,16 +189,18 @@ def test_splunk_logger_unhandled_failure(self):
         )
 
         with (  # noqa: E999
+            patch("common.log_firehose.firehose_client") as mock_firehose_client,  # noqa: E999
+            patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
             patch("logging_decorator.datetime") as mock_datetime,  # noqa: E999
-            patch("logging_decorator.time") as mock_time,  # noqa: E999
-            patch("logging_decorator.logger") as mock_logger,  # noqa: E999
-            patch("logging_decorator.firehose_client") as mock_firehose_client,  # noqa: E999
+            patch("logging_decorator.time") as mock_start_time,  # noqa: E999
+            patch("common.log_decorator.time") as mock_end_time,  # noqa: E999
             patch(
                 "file_level_validation.validate_content_headers",
                 side_effect=ValueError("Test exception"),
             ),  # noqa: E999
         ):  # noqa: E999
-            mock_time.time.side_effect = [1672531200, 1672531200.123456]
+            mock_start_time.time.return_value = 1672531200
+            mock_end_time.time.return_value = 1672531200.123456
             mock_datetime.now.return_value = datetime(2024, 1, 1, 12, 0, 0)
             with self.assertRaises(ValueError):
                 file_level_validation(deepcopy(MOCK_FILE_DETAILS.event_full_permissions_dict))

lambdas/recordprocessor/tests/test_recordprocessor_edge_cases.py

@@ -4,7 +4,7 @@
 from unittest.mock import call, patch
 
 from batch_processor import process_csv_to_fhir
-from tests.utils_for_recordprocessor_tests.utils_for_recordprocessor_tests import (
+from utils_for_recordprocessor_tests.utils_for_recordprocessor_tests import (
     create_patch,
 )
 
@@ -21,7 +21,7 @@ def setUp(self):
         self.mock_make_and_move = create_patch("file_level_validation.make_and_upload_ack_file")
         self.mock_move_file = create_patch("file_level_validation.move_file")
         self.mock_get_permitted_operations = create_patch("file_level_validation.get_permitted_operations")
-        self.mock_firehose_client = create_patch("logging_decorator.firehose_client")
+        self.mock_firehose_client = create_patch("common.log_firehose.firehose_client")
         self.mock_update_audit_table_status = create_patch("batch_processor.update_audit_table_status")
 
     def tearDown(self):

lambdas/recordprocessor/tests/test_recordprocessor_main.py

@@ -13,6 +13,7 @@
 from utils_for_recordprocessor_tests.mock_environment_variables import (
     MOCK_ENVIRONMENT_DICT,
     BucketNames,
+    Firehose,
     Kinesis,
 )
 from utils_for_recordprocessor_tests.utils_for_recordprocessor_tests import (
@@ -437,15 +438,17 @@ def test_e2e_kinesis_failed(self):
         kinesis_client.delete_stream(StreamName=Kinesis.STREAM_NAME, EnforceConsumerDeletion=True)
 
         with (  # noqa: E999
-            patch("logging_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
+            patch("common.log_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
             patch("logging_decorator.datetime") as mock_datetime,  # noqa: E999
-            patch("logging_decorator.time") as mock_time,  # noqa: E999
+            patch("logging_decorator.time") as mock_start_time,  # noqa: E999
+            patch("common.log_decorator.time") as mock_end_time,  # noqa: E999
         ):  # noqa: E999
-            mock_time.time.side_effect = [1672531200, 1672531200.123456]
+            mock_start_time.time.return_value = 1672531200
+            mock_end_time.time.return_value = 1672531200.123456
             mock_datetime.now.return_value = datetime(2024, 1, 1, 12, 0, 0)
             main(test_file.event_full_permissions)
 
-        # Since the failure occured at row level, not file level, the ack file should still be created
+        # Since the failure occurred at row level, not file level, the ack file should still be created
         # and firehose logs should indicate a successful file level validation
         table_entry = dynamo_db_client.get_item(
             TableName=AUDIT_TABLE_NAME,
@@ -463,7 +466,7 @@ def test_e2e_kinesis_failed(self):
             "statusCode": 200,
             "message": "Successfully sent for record processing",
         }
-        mock_send_log_to_firehose.assert_called_with(expected_log_data)
+        mock_send_log_to_firehose.assert_called_with(Firehose.STREAM_NAME, expected_log_data)
         self.assertDictEqual(
             table_entry,
             {

lambdas/recordprocessor/tests/utils_for_recordprocessor_tests/mock_environment_variables.py

@@ -37,6 +37,7 @@ class Sqs:
     "SOURCE_BUCKET_NAME": BucketNames.SOURCE,
     "ACK_BUCKET_NAME": BucketNames.DESTINATION,
     "SHORT_QUEUE_PREFIX": "imms-batch-internal-dev",
+    "SPLUNK_FIREHOSE_NAME": Firehose.STREAM_NAME,
     "KINESIS_STREAM_NAME": Kinesis.STREAM_NAME,
     "KINESIS_STREAM_ARN": f"arn:aws:kinesis:{REGION_NAME}:123456789012:stream/{Kinesis.STREAM_NAME}",
     "FIREHOSE_STREAM_NAME": Firehose.STREAM_NAME,

lambdas/shared/src/common/log_decorator.py

@@ -14,8 +14,8 @@
 
 
 def generate_and_send_logs(
-    stream_name,
-    start_time,
+    stream_name: str,
+    start_time: float,
     base_log_data: dict,
     additional_log_data: dict,
     use_ms_precision: bool = False,

lambdas/shared/src/common/log_firehose.py

@@ -6,7 +6,7 @@
 )
 
 
-def send_log_to_firehose(stream_name, log_data: dict) -> None:
+def send_log_to_firehose(stream_name: str, log_data: dict) -> None:
     """Sends the log_message to Firehose"""
     try:
         record = {"Data": json.dumps({"event": log_data}).encode("utf-8")}