VED-1010: refactor audit table functions (#1137)

* Refactor upsert_audit_table to create_audit_table_item

* Refactor update_audit_table_status and set_audit_table_ingestion_start_time in recordprocessor to use common update_audit_table_item

* Refactor ack_backend to use shared update

Also moved other audit table functions to shared file

* Use shared function in batch_processor_filter

* Refactor update_audit_table_item

* Consolidate unit tests

* Tidy up create_audit_table_item and unit tests

* Delete consolidated unit test files

* Self review

* Tidy up after rebase

* Self review 2

* Fix linting issue

* Address review feedback

* Tidy up error handling

Author: edhall-nhs

lambdas/ack_backend/src/ack_processor.py

@@ -2,7 +2,7 @@
 
 import json
 
-from audit_table import increment_records_failed_count
+from common.batch.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

lambdas/ack_backend/src/audit_table.py

@@ -7,15 +7,11 @@
 
 from botocore.exceptions import ClientError
 
-from audit_table import (
-    change_audit_table_status_to_processed,
-    get_record_count_and_failures_by_message_id,
-    set_audit_record_success_count_and_end_time,
-)
 from common.aws_s3_utils import move_file
+from common.batch.audit_table import get_record_count_and_failures_by_message_id, update_audit_table_item
 from common.clients import get_s3_client, logger
 from common.log_decorator import generate_and_send_logs
-from common.models.batch_constants import ACK_BUCKET_NAME, SOURCE_BUCKET_NAME
+from common.models.batch_constants import ACK_BUCKET_NAME, SOURCE_BUCKET_NAME, AuditTableKeys, FileStatus
 from constants import (
     ACK_HEADERS,
     BATCH_FILE_ARCHIVE_DIR,
@@ -81,12 +77,21 @@ def complete_batch_file_process(
     move_file(SOURCE_BUCKET_NAME, f"{BATCH_FILE_PROCESSING_DIR}/{file_key}", f"{BATCH_FILE_ARCHIVE_DIR}/{file_key}")
 
     total_ack_rows_processed, total_failures = get_record_count_and_failures_by_message_id(message_id)
-    change_audit_table_status_to_processed(file_key, message_id)
+    update_audit_table_item(
+        file_key=file_key, message_id=message_id, attrs_to_update={AuditTableKeys.STATUS: FileStatus.PROCESSED}
+    )
 
     # Consider creating time utils and using datetime instead of time
     ingestion_end_time = time.strftime("%Y%m%dT%H%M%S00", time.gmtime())
     successful_record_count = total_ack_rows_processed - total_failures
-    set_audit_record_success_count_and_end_time(file_key, message_id, successful_record_count, ingestion_end_time)
+    update_audit_table_item(
+        file_key=file_key,
+        message_id=message_id,
+        attrs_to_update={
+            AuditTableKeys.RECORDS_SUCCEEDED: successful_record_count,
+            AuditTableKeys.INGESTION_END_TIME: ingestion_end_time,
+        },
+    )
 
     result = {
         "message_id": message_id,

lambdas/ack_backend/src/update_ack_file.py

@@ -43,7 +43,7 @@
 
 
 @patch.dict(os.environ, MOCK_ENVIRONMENT_DICT)
-@patch("audit_table.AUDIT_TABLE_NAME", AUDIT_TABLE_NAME)
+@patch("common.batch.audit_table.AUDIT_TABLE_NAME", AUDIT_TABLE_NAME)
 @mock_aws
 class TestAckProcessor(unittest.TestCase):
     """Tests for the ack processor lambda handler."""

lambdas/ack_backend/tests/test_ack_processor.py

@@ -386,8 +386,8 @@ def test_splunk_update_ack_file_not_logged(self):
             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(
-                "update_ack_file.change_audit_table_status_to_processed"
-            ) as mock_change_audit_table_status_to_processed,  # noqa: E999
+                "update_ack_file.update_audit_table_item",
+            ) as mock_update_audit_table_item,  # noqa: E999
             patch("ack_processor.increment_records_failed_count"),  # noqa: E999
         ):  # noqa: E999
             result = lambda_handler(generate_event(messages), context={})
@@ -413,7 +413,7 @@ def test_splunk_update_ack_file_not_logged(self):
                 call(self.stream_name, last_logger_info_call_args),
             ]
         )
-        mock_change_audit_table_status_to_processed.assert_not_called()
+        mock_update_audit_table_item.assert_not_called()
 
     def test_splunk_update_ack_file_logged(self):
         """Tests that update_ack_file is logged if we have sent acks for the whole file"""
@@ -428,12 +428,7 @@ def test_splunk_update_ack_file_logged(self):
             patch("update_ack_file.datetime") as mock_datetime,
             patch("common.log_decorator.logger") as mock_logger,  # noqa: E999
             patch("update_ack_file.get_record_count_and_failures_by_message_id", return_value=(99, 2)),
-            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_audit_record_success_count_and_end_time"
-            ) as mock_set_records_succeeded_count_and_end_time,  # noqa: E999
+            patch("update_ack_file.update_audit_table_item") as mock_update_audit_table_item,  # noqa: E999
             patch("ack_processor.increment_records_failed_count"),  # noqa: E999
         ):  # noqa: E999
             mock_datetime.now.return_value = ValidValues.fixed_datetime
@@ -469,8 +464,8 @@ def test_splunk_update_ack_file_logged(self):
                 call(self.stream_name, last_logger_info_call_args),
             ]
         )
-        mock_change_audit_table_status_to_processed.assert_called()
-        mock_set_records_succeeded_count_and_end_time.assert_called()
+
+        self.assertEqual(mock_update_audit_table_item.call_count, 2)
 
 
 if __name__ == "__main__":

lambdas/ack_backend/tests/test_audit_table.py

@@ -33,18 +33,15 @@ def setUp(self):
         self.logger_patcher = patch("update_ack_file.logger")
         self.mock_logger = self.logger_patcher.start()
 
-        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.update_audit_table_item_patcher = patch("update_ack_file.update_audit_table_item")
+        self.mock_update_audit_table_item = self.update_audit_table_item_patcher.start()
         self.get_record_and_failure_count_patcher = patch("update_ack_file.get_record_count_and_failures_by_message_id")
         self.mock_get_record_and_failure_count = self.get_record_and_failure_count_patcher.start()
-        self.set_records_succeeded_count_patcher = patch("update_ack_file.set_audit_record_success_count_and_end_time")
-        self.mock_set_records_succeeded_count = self.set_records_succeeded_count_patcher.start()
 
     def tearDown(self):
         self.logger_patcher.stop()
-        self.change_audit_status_patcher.stop()
+        self.update_audit_table_item_patcher.stop()
         self.get_record_and_failure_count_patcher.stop()
-        self.set_records_succeeded_count_patcher.stop()
 
     def test_audit_table_updated_correctly_when_ack_process_complete(self):
         """VED-167 - Test that the audit table has been updated correctly"""
@@ -73,5 +70,4 @@ def test_audit_table_updated_correctly_when_ack_process_complete(self):
 
         # Assert: Only check audit table interactions
         self.mock_get_record_and_failure_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()
+        self.assertEqual(self.mock_update_audit_table_item.call_count, 2)