unit tests

Author: JamesW1-NHS

lambdas/ack_backend/src/ack_processor.py

@@ -2,7 +2,7 @@
 
 import json
 
-# from audit_table import increment_record_counter
+from audit_table import increment_records_failed_count
 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
@@ -51,9 +51,7 @@ def lambda_handler(event, _):
 
             ack_data_rows.append(convert_message_to_ack_row(message, created_at_formatted_string))
             if message.get("diagnostics"):
-                # TODO: unit tests will fail; we need to mock the audit table to get this to work.
-                print("TODO: increment the record counter")
-                # increment_record_counter(message_id)
+                increment_records_failed_count(message_id)
 
     update_ack_file(file_key, created_at_formatted_string, ack_data_rows)
 

lambdas/ack_backend/src/update_ack_file.py

@@ -4,7 +4,11 @@
 
 from botocore.exceptions import ClientError
 
-from audit_table import change_audit_table_status_to_processed, get_record_count_by_message_id
+from audit_table import (
+    change_audit_table_status_to_processed,
+    get_record_count_by_message_id,
+    set_records_succeeded_count,
+)
 from common.aws_s3_utils import move_file
 from common.clients import get_s3_client, logger
 from constants import (
@@ -73,9 +77,7 @@ def complete_batch_file_process(
 
     total_ack_rows_processed = get_record_count_by_message_id(message_id)
     change_audit_table_status_to_processed(file_key, message_id)
-    # TODO: we need to mock this here to pass unit tests
-    print("TODO: Call set_record_success_count")
-    # set_record_success_count(message_id)
+    set_records_succeeded_count(message_id)
 
     return {
         "message_id": message_id,

lambdas/ack_backend/tests/test_ack_processor.py

@@ -111,6 +111,19 @@ def assert_audit_entry_status_equals(self, message_id: str, status: str) -> None
         actual_status = audit_entry.get("status", {}).get("S")
         self.assertEqual(actual_status, status)
 
+    def assert_audit_entry_counts_equal(self, message_id: str, expected_counts: dict) -> None:
+        """Checks the audit entry counts are as expected"""
+        audit_entry = self.dynamodb_client.get_item(
+            TableName=AUDIT_TABLE_NAME, Key={"message_id": {"S": message_id}}
+        ).get("Item")
+
+        actual_counts = {}
+        actual_counts["record_count"] = audit_entry.get("record_count", {}).get("N")
+        actual_counts["records_succeeded"] = audit_entry.get("records_succeeded", {}).get("N")
+        actual_counts["records_failed"] = audit_entry.get("records_failed", {}).get("N")
+
+        self.assertEqual(actual_counts, expected_counts)
+
     def test_lambda_handler_main_multiple_records(self):
         """Test lambda handler with multiple records."""
         # Set up an audit entry which does not yet have record_count recorded
@@ -164,6 +177,11 @@ def test_lambda_handler_main_multiple_records(self):
                 {"body": json.dumps(array_of_mixed_success_and_failure_messages)},
             ]
         }
+        expected_entry_counts = {
+            "record_count": None,
+            "records_succeeded": None,
+            "records_failed": "6",
+        }
 
         response = lambda_handler(event=event, context={})
 
@@ -177,6 +195,7 @@ def test_lambda_handler_main_multiple_records(self):
             ],
             existing_file_content=ValidValues.ack_headers,
         )
+        self.assert_audit_entry_counts_equal("row", expected_entry_counts)
 
     def test_lambda_handler_main(self):
         """Test lambda handler with consitent ack_file_name and message_template."""
@@ -217,12 +236,20 @@ def test_lambda_handler_main(self):
                 "messages": [{"row_id": "row^1", "diagnostics": "SHOULD BE A DICTIONARY, NOT A STRING"}],
             },
         ]
+        expected_records_failed = [None, "3", "8", "8", "9"]
+        expected_entry_counts = {
+            "record_count": None,
+            "records_succeeded": None,
+            "records_failed": None,
+        }
 
-        for test_case in test_cases:
+        for test_case, expected_failures in zip(test_cases, expected_records_failed):
+            expected_entry_counts["records_failed"] = expected_failures
             # Test scenario where there is no existing ack file
             with self.subTest(msg=f"No existing ack file: {test_case['description']}"):
                 response = lambda_handler(event=self.generate_event(test_case["messages"]), context={})
                 self.assertEqual(response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
+                self.assert_audit_entry_counts_equal("row", expected_entry_counts)
                 validate_ack_file_content(self.s3_client, test_case["messages"])
 
                 self.s3_client.delete_object(
@@ -252,6 +279,11 @@ def test_lambda_handler_updates_ack_file_but_does_not_mark_complete_when_records
             for i in range(1, 4)
         ]
         test_event = {"Records": [{"body": json.dumps(array_of_success_messages)}]}
+        expected_entry_counts = {
+            "record_count": "100",
+            "records_succeeded": None,
+            "records_failed": None,
+        }
 
         response = lambda_handler(event=test_event, context={})
 
@@ -268,6 +300,7 @@ def test_lambda_handler_updates_ack_file_but_does_not_mark_complete_when_records
             is_complete=False,
         )
         self.assert_audit_entry_status_equals(mock_batch_message_id, "Preprocessed")
+        self.assert_audit_entry_counts_equal(mock_batch_message_id, expected_entry_counts)
 
     def test_lambda_handler_updates_ack_file_and_marks_complete_when_all_records_processed(self):
         """
@@ -305,6 +338,11 @@ def test_lambda_handler_updates_ack_file_and_marks_complete_when_all_records_pro
         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)}]}
+        expected_entry_counts = {
+            "record_count": "100",
+            "records_succeeded": "100",
+            "records_failed": None,
+        }
 
         response = lambda_handler(event=test_event, context={})
 
@@ -319,6 +357,61 @@ def test_lambda_handler_updates_ack_file_and_marks_complete_when_all_records_pro
             is_complete=True,
         )
         self.assert_audit_entry_status_equals(mock_batch_message_id, "Processed")
+        self.assert_audit_entry_counts_equal(mock_batch_message_id, expected_entry_counts)
+
+    def test_lambda_handler_sets_records_succeeded(self):
+        """
+        Test that the records_succeeded count is set when all records have been processed.
+        """
+        mock_batch_message_id = "75db20e6-c0b5-4012-a8bc-f861a1dd4b22"
+
+        # Original source file had 100 records
+        add_audit_entry_to_table(self.dynamodb_client, mock_batch_message_id, record_count=100)
+
+        # Previous invocations have already created and added to the temp ack file
+        existing_ack_content = generate_sample_existing_ack_content()
+
+        self.s3_client.put_object(
+            Bucket=BucketNames.DESTINATION,
+            Key=MOCK_MESSAGE_DETAILS.temp_ack_file_key,
+            Body=StringIO(existing_ack_content).getvalue(),
+        )
+
+        array_of_success_messages = [
+            {
+                **BASE_SUCCESS_MESSAGE,
+                "row_id": f"{mock_batch_message_id}^{i}",
+                "imms_id": f"imms_{i}",
+                "local_id": f"local^{i}",
+            }
+            for i in range(50, 75)
+        ]
+        array_of_failure_messages = [
+            {
+                **BASE_FAILURE_MESSAGE,
+                "row_id": f"{mock_batch_message_id}^{i}",
+                "local_id": f"local^{i}",
+                "diagnostics": DiagnosticsDictionaries.UNHANDLED_ERROR,
+            }
+            for i in range(75, 100)
+        ]
+
+        # Include the EoF message in the event
+        array_of_messages = deepcopy(array_of_success_messages) + deepcopy(array_of_failure_messages)
+        all_messages_plus_eof = array_of_messages
+        all_messages_plus_eof.append(MOCK_MESSAGE_DETAILS.eof_message)
+        test_event = {"Records": [{"body": json.dumps(all_messages_plus_eof)}]}
+        expected_entry_counts = {
+            "record_count": "100",
+            "records_succeeded": "75",
+            "records_failed": "25",
+        }
+
+        response = lambda_handler(event=test_event, context={})
+
+        self.assertEqual(response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
+        self.assert_audit_entry_status_equals(mock_batch_message_id, "Processed")
+        self.assert_audit_entry_counts_equal(mock_batch_message_id, expected_entry_counts)
 
     def test_lambda_handler_error_scenarios(self):
         """Test that the lambda handler raises appropriate exceptions for malformed event data."""

lambdas/ack_backend/tests/test_splunk_logging.py

@@ -224,6 +224,7 @@ def test_splunk_logging_statuscode_diagnostics(
             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.increment_records_failed_count") as mock_increment_records_failed_count,  # noqa: E999
             ):  # noqa: E999
                 result = lambda_handler(
                     event=generate_event([{"diagnostics": test_case["diagnostics"]}]),
@@ -253,6 +254,7 @@ def test_splunk_logging_statuscode_diagnostics(
                     call(self.stream_name, expected_second_logger_info_data),
                 ]
             )
+            mock_increment_records_failed_count.assert_called()
 
     def test_splunk_logging_multiple_rows(self):
         """Tests logging for multiple objects in the body of the event"""
@@ -311,6 +313,7 @@ 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("common.log_decorator.logger") as mock_logger,  # noqa: E999
+            patch("ack_processor.increment_records_failed_count") as mock_increment_records_failed_count,  # noqa: E999
         ):  # noqa: E999
             result = lambda_handler(generate_event(messages), context={})
 
@@ -363,6 +366,7 @@ def test_splunk_logging_multiple_with_diagnostics(
                 call(self.stream_name, expected_fourth_logger_info_data),
             ]
         )
+        mock_increment_records_failed_count.assert_called()
 
     def test_splunk_update_ack_file_not_logged(self):
         """Tests that update_ack_file is not logged if we have sent acks for less than the whole file"""
@@ -419,6 +423,7 @@ def test_splunk_update_ack_file_logged(self):
             patch(
                 "update_ack_file.change_audit_table_status_to_processed"
             ) as mock_change_audit_table_status_to_processed,  # noqa: E999
+            patch("update_ack_file.set_records_succeeded_count") as mock_set_records_succeeded_count,  # noqa: E999
         ):  # noqa: E999
             result = lambda_handler(generate_event(messages, include_eof_message=True), context={})
 
@@ -453,6 +458,7 @@ def test_splunk_update_ack_file_logged(self):
             ]
         )
         mock_change_audit_table_status_to_processed.assert_called()
+        mock_set_records_succeeded_count.assert_called()
 
 
 if __name__ == "__main__":

lambdas/ack_backend/tests/test_update_ack_file_flow.py

@@ -39,6 +39,8 @@ def setUp(self):
         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()
+        self.set_records_succeeded_count_patcher = patch("update_ack_file.set_records_succeeded_count")
+        self.mock_set_records_succeeded_count = self.set_records_succeeded_count_patcher.start()
 
     def tearDown(self):
         self.logger_patcher.stop()
@@ -72,3 +74,4 @@ def test_audit_table_updated_correctly_when_ack_process_complete(self):
         # 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)
+        self.mock_set_records_succeeded_count.assert_called_once()