VED-924 (Also contains VED-925 and VED-926) Modify ack file creation and completion process (#1023)

Author: dlzhry2nhs

lambdas/ack_backend/src/ack_processor.py

@@ -2,10 +2,10 @@
 
 import json
 
+from common.batch.eof_utils import is_eof_message
 from convert_message_to_ack_row import convert_message_to_ack_row
 from logging_decorators import ack_lambda_handler_logging_decorator
 from update_ack_file import complete_batch_file_process, update_ack_file
-from utils_for_ack_lambda import is_ack_processing_complete
 
 
 @ack_lambda_handler_logging_decorator
@@ -22,9 +22,11 @@ def lambda_handler(event, _):
     file_key = None
     created_at_formatted_string = None
     message_id = None
+    supplier = None
+    vaccine_type = None
 
     ack_data_rows = []
-    total_ack_rows_processed = 0
+    file_processing_complete = False
 
     for i, record in enumerate(event["Records"]):
         try:
@@ -42,18 +44,16 @@ def lambda_handler(event, _):
             created_at_formatted_string = incoming_message_body[0].get("created_at_formatted_string")
 
         for message in incoming_message_body:
+            if is_eof_message(message):
+                file_processing_complete = True
+                break
+
             ack_data_rows.append(convert_message_to_ack_row(message, created_at_formatted_string))
 
     update_ack_file(file_key, created_at_formatted_string, ack_data_rows)
 
-    # Get the row count of the final processed record
-    # Format of the row id is {batch_message_id}^{row_number}
-    total_ack_rows_processed = int(incoming_message_body[-1].get("row_id", "").split("^")[1])
-
-    if is_ack_processing_complete(message_id, total_ack_rows_processed):
-        complete_batch_file_process(
-            message_id, supplier, vaccine_type, created_at_formatted_string, file_key, total_ack_rows_processed
-        )
+    if file_processing_complete:
+        complete_batch_file_process(message_id, supplier, vaccine_type, created_at_formatted_string, file_key)
 
     return {
         "statusCode": 200,

lambdas/ack_backend/src/update_ack_file.py

@@ -4,7 +4,7 @@
 
 from botocore.exceptions import ClientError
 
-from audit_table import change_audit_table_status_to_processed
+from audit_table import change_audit_table_status_to_processed, get_record_count_by_message_id
 from common.aws_s3_utils import move_file
 from common.clients import get_s3_client, logger
 from constants import (
@@ -61,7 +61,6 @@ def complete_batch_file_process(
     vaccine_type: str,
     created_at_formatted_string: str,
     file_key: str,
-    total_ack_rows_processed: int,
 ) -> dict:
     """Mark the batch file as processed. This involves moving the ack and original file to destinations and updating
     the audit table status"""
@@ -72,6 +71,7 @@ def complete_batch_file_process(
         get_source_bucket_name(), f"{BATCH_FILE_PROCESSING_DIR}/{file_key}", f"{BATCH_FILE_ARCHIVE_DIR}/{file_key}"
     )
 
+    total_ack_rows_processed = get_record_count_by_message_id(message_id)
     change_audit_table_status_to_processed(file_key, message_id)
 
     return {
@@ -111,7 +111,7 @@ def update_ack_file(
     """Updates the ack file with the new data row based on the given arguments"""
     ack_filename = f"{file_key.replace('.csv', f'_BusAck_{created_at_formatted_string}.csv')}"
     temp_ack_file_key = f"{TEMP_ACK_DIR}/{ack_filename}"
-    archive_ack_file_key = f"{COMPLETED_ACK_DIR}/{ack_filename}"
+    completed_ack_file_key = f"{COMPLETED_ACK_DIR}/{ack_filename}"
     accumulated_csv_content = obtain_current_ack_content(temp_ack_file_key)
 
     for row in ack_data_rows:
@@ -123,4 +123,4 @@ def update_ack_file(
     ack_bucket_name = get_ack_bucket_name()
 
     get_s3_client().upload_fileobj(csv_file_like_object, ack_bucket_name, temp_ack_file_key)
-    logger.info("Ack file updated to %s: %s", ack_bucket_name, archive_ack_file_key)
+    logger.info("Ack file updated to %s: %s", ack_bucket_name, completed_ack_file_key)

lambdas/ack_backend/src/utils_for_ack_lambda.py

@@ -3,8 +3,9 @@
 import json
 import os
 import unittest
+from copy import deepcopy
 from io import StringIO
-from unittest.mock import Mock, patch
+from unittest.mock import patch
 
 from boto3 import client as boto3_client
 from moto import mock_aws
@@ -30,7 +31,6 @@
     DiagnosticsDictionaries,
     ValidValues,
 )
-from utils_for_ack_lambda import _BATCH_EVENT_ID_TO_RECORD_COUNT_MAP
 
 with patch.dict("os.environ", MOCK_ENVIRONMENT_DICT):
     from ack_processor import lambda_handler
@@ -269,45 +269,6 @@ def test_lambda_handler_updates_ack_file_but_does_not_mark_complete_when_records
         )
         self.assert_audit_entry_status_equals(mock_batch_message_id, "Preprocessed")
 
-    @patch("utils_for_ack_lambda.get_record_count_by_message_id", return_value=500)
-    def test_lambda_handler_uses_message_id_to_record_count_cache_to_reduce_ddb_calls(self, mock_get_record_count: Mock):
-        """The DynamoDB Audit table is used to store the total record count for each source file. To reduce calls each
-        time - this test checks that we cache the value as this lambda is called many times for large files"""
-        mock_batch_message_id = "622cdeea-461e-4a83-acb5-7871d47ddbcd"
-
-        # Original source file had 500 records
-        add_audit_entry_to_table(self.dynamodb_client, mock_batch_message_id, record_count=500)
-
-        message_one = [
-            {**BASE_SUCCESS_MESSAGE, "row_id": f"{mock_batch_message_id}^1", "imms_id": "imms_1", "local_id": "local^1"}
-        ]
-        message_two = [
-            {**BASE_SUCCESS_MESSAGE, "row_id": f"{mock_batch_message_id}^2", "imms_id": "imms_2", "local_id": "local^2"}
-        ]
-        test_event_one = {"Records": [{"body": json.dumps(message_one)}]}
-        test_event_two = {"Records": [{"body": json.dumps(message_two)}]}
-
-        response = lambda_handler(event=test_event_one, context={})
-        self.assertEqual(response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
-        second_invocation_response = lambda_handler(event=test_event_two, context={})
-        self.assertEqual(second_invocation_response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
-
-        # Assert that the DDB call is only performed once on the first invocation
-        mock_get_record_count.assert_called_once_with(mock_batch_message_id)
-        validate_ack_file_content(
-            self.s3_client,
-            [*message_one, *message_two],
-            existing_file_content=ValidValues.ack_headers,
-        )
-        self.assert_ack_and_source_file_locations_correct(
-            MOCK_MESSAGE_DETAILS.file_key,
-            MOCK_MESSAGE_DETAILS.temp_ack_file_key,
-            MOCK_MESSAGE_DETAILS.archive_ack_file_key,
-            is_complete=False,
-        )
-        self.assertEqual(_BATCH_EVENT_ID_TO_RECORD_COUNT_MAP[mock_batch_message_id], 500)
-        self.assert_audit_entry_status_equals(mock_batch_message_id, "Preprocessed")
-
     def test_lambda_handler_updates_ack_file_and_marks_complete_when_all_records_processed(self):
         """
         Test that the batch file process is marked as complete when all records have been processed.
@@ -339,7 +300,11 @@ def test_lambda_handler_updates_ack_file_and_marks_complete_when_all_records_pro
             }
             for i in range(50, 101)
         ]
-        test_event = {"Records": [{"body": json.dumps(array_of_success_messages)}]}
+
+        # Include the EoF message in the event
+        all_messages_plus_eof = deepcopy(array_of_success_messages)
+        all_messages_plus_eof.append(MOCK_MESSAGE_DETAILS.eof_message)
+        test_event = {"Records": [{"body": json.dumps(all_messages_plus_eof)}]}
 
         response = lambda_handler(event=test_event, context={})
 

lambdas/ack_backend/tests/test_ack_processor.py

@@ -9,13 +9,13 @@
 from boto3 import client as boto3_client
 from moto import mock_aws
 
-from tests.utils.generic_setup_and_teardown_for_ack_backend import (
+from utils.generic_setup_and_teardown_for_ack_backend import (
     GenericSetUp,
     GenericTearDown,
 )
-from tests.utils.mock_environment_variables import MOCK_ENVIRONMENT_DICT, BucketNames
-from tests.utils.utils_for_ack_backend_tests import generate_event
-from tests.utils.values_for_ack_backend_tests import (
+from utils.mock_environment_variables import MOCK_ENVIRONMENT_DICT, BucketNames
+from utils.utils_for_ack_backend_tests import generate_event
+from utils.values_for_ack_backend_tests import (
     EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS,
     DiagnosticsDictionaries,
     InvalidValues,
@@ -120,7 +120,6 @@ def test_splunk_logging_successful_rows(self):
         for operation in ["CREATE", "UPDATE", "DELETE"]:
             with (  # noqa: E999
                 patch("common.log_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
-                patch("ack_processor.is_ack_processing_complete", return_value=False),
                 patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
             ):  # noqa: E999
                 result = lambda_handler(
@@ -224,7 +223,6 @@ def test_splunk_logging_statuscode_diagnostics(
         for test_case in test_cases:
             with (  # noqa: E999
                 patch("common.log_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
-                patch("ack_processor.is_ack_processing_complete", return_value=False),
                 patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
             ):  # noqa: E999
                 result = lambda_handler(
@@ -262,7 +260,6 @@ def test_splunk_logging_multiple_rows(self):
 
         with (  # noqa: E999
             patch("common.log_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
-            patch("ack_processor.is_ack_processing_complete", return_value=False),
             patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
         ):  # noqa: E999
             result = lambda_handler(generate_event(messages), context={})
@@ -313,7 +310,6 @@ def test_splunk_logging_multiple_with_diagnostics(
 
         with (  # noqa: E999
             patch("common.log_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
-            patch("ack_processor.is_ack_processing_complete", return_value=False),
             patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
         ):  # noqa: E999
             result = lambda_handler(generate_event(messages), context={})
@@ -369,7 +365,6 @@ def test_splunk_logging_multiple_with_diagnostics(
         )
 
     def test_splunk_update_ack_file_not_logged(self):
-        self.maxDiff = None
         """Tests that update_ack_file is not logged if we have sent acks for less than the whole file"""
         # send 98 messages
         messages = []
@@ -380,7 +375,6 @@ def test_splunk_update_ack_file_not_logged(self):
         with (  # noqa: E999
             patch("common.log_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
             patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
-            patch("ack_processor.is_ack_processing_complete", return_value=False),
             patch(
                 "update_ack_file.change_audit_table_status_to_processed"
             ) as mock_change_audit_table_status_to_processed,  # noqa: E999
@@ -421,12 +415,12 @@ def test_splunk_update_ack_file_logged(self):
         with (  # noqa: E999
             patch("common.log_decorator.send_log_to_firehose") as mock_send_log_to_firehose,  # noqa: E999
             patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
-            patch("ack_processor.is_ack_processing_complete", return_value=True),
+            patch("update_ack_file.get_record_count_by_message_id", return_value=99),
             patch(
                 "update_ack_file.change_audit_table_status_to_processed"
             ) as mock_change_audit_table_status_to_processed,  # noqa: E999
         ):  # noqa: E999
-            result = lambda_handler(generate_event(messages), context={})
+            result = lambda_handler(generate_event(messages, include_eof_message=True), context={})
 
         self.assertEqual(result, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
 

lambdas/ack_backend/tests/test_splunk_logging.py

@@ -37,6 +37,8 @@ def setUp(self):
 
         self.change_audit_status_patcher = patch("update_ack_file.change_audit_table_status_to_processed")
         self.mock_change_audit_status = self.change_audit_status_patcher.start()
+        self.get_record_count_patcher = patch("update_ack_file.get_record_count_by_message_id")
+        self.mock_get_record_count = self.get_record_count_patcher.start()
 
     def tearDown(self):
         self.logger_patcher.stop()
@@ -56,6 +58,7 @@ def test_audit_table_updated_correctly_when_ack_process_complete(self):
         self.s3_client.put_object(
             Bucket=self.ack_bucket_name, Key=f"TempAck/audit_table_test_BusAck_{mock_created_at_string}.csv"
         )
+        self.mock_get_record_count.return_value = 10
 
         # Act
         update_ack_file.complete_batch_file_process(
@@ -64,8 +67,8 @@ def test_audit_table_updated_correctly_when_ack_process_complete(self):
             vaccine_type="vaccine-type",
             created_at_formatted_string=mock_created_at_string,
             file_key=file_key,
-            total_ack_rows_processed=3,
         )
 
-        # Assert: Only check audit table update
+        # Assert: Only check audit table interactions
+        self.mock_get_record_count.assert_called_once_with(message_id)
         self.mock_change_audit_status.assert_called_once_with(file_key, message_id)

lambdas/ack_backend/tests/test_update_ack_file_flow.py

@@ -21,7 +21,7 @@ def add_audit_entry_to_table(dynamodb_client, batch_event_message_id: str, recor
     dynamodb_client.put_item(TableName=AUDIT_TABLE_NAME, Item=audit_table_entry)
 
 
-def generate_event(test_messages: list[dict]) -> dict:
+def generate_event(test_messages: list[dict], include_eof_message: bool = False) -> dict:
     """
     Returns an event where each message in the incoming message body list is based on a standard mock message,
     updated with the details from the corresponsing message in the given test_messages list.
@@ -34,6 +34,10 @@ def generate_event(test_messages: list[dict]) -> dict:
         )
         for message in test_messages
     ]
+
+    if include_eof_message:
+        incoming_message_body.append(MOCK_MESSAGE_DETAILS.eof_message)
+
     return {"Records": [{"body": json.dumps(incoming_message_body)}]}
 
 

lambdas/ack_backend/tests/utils/utils_for_ack_backend_tests.py

@@ -139,6 +139,15 @@ def __init__(
             "diagnostics": DiagnosticsDictionaries.NO_PERMISSIONS,
         }
 
+        self.eof_message = {
+            "created_at_formatted_string": self.created_at_formatted_string,
+            "file_key": self.file_key,
+            "message": "EOF",
+            "row_id": self.row_id,
+            "supplier": self.supplier,
+            "vax_type": self.vaccine_type,
+        }
+
 
 class MockMessageDetails:
     """Class containing mock message details for use in tests"""