Merge branch 'master' into VED-124-Batch-Splunk-retool

Author: JamesW1-NHS

ack_backend/Makefile

@@ -6,6 +6,6 @@ package: build
 	docker run --rm -v $(shell pwd)/build:/build ack-lambda-build
 
 test:
-	python -m unittest
+	@PYTHONPATH=src:tests python -m unittest
 
 .PHONY: build package 

ack_backend/poetry.lock

@@ -5,13 +5,18 @@
 
 REGION_NAME = "eu-west-2"
 
-s3_client = boto3_client("s3", region_name=REGION_NAME)
 firehose_client = boto3_client("firehose", region_name=REGION_NAME)
 lambda_client = boto3_client('lambda', region_name=REGION_NAME)
 dynamodb_client = boto3_client("dynamodb", region_name=REGION_NAME)
 
 dynamodb_resource = boto3_resource("dynamodb", region_name=REGION_NAME)
 
+s3_client = None
+def get_s3_client():
+    global s3_client
+    if s3_client is None:
+        s3_client = boto3_client("s3", region_name=REGION_NAME)
+    return s3_client
 
 # Logger
 logging.basicConfig(level="INFO")

ack_backend/src/clients.py

@@ -2,13 +2,18 @@
 
 import os
 
-SOURCE_BUCKET_NAME = os.getenv("SOURCE_BUCKET_NAME")
-ACK_BUCKET_NAME = os.getenv("ACK_BUCKET_NAME")
 AUDIT_TABLE_NAME = os.getenv("AUDIT_TABLE_NAME")
+FILE_NAME_PROC_LAMBDA_NAME = os.getenv("FILE_NAME_PROC_LAMBDA_NAME")
 AUDIT_TABLE_FILENAME_GSI = "filename_index"
 AUDIT_TABLE_QUEUE_NAME_GSI = "queue_name_index"
-FILE_NAME_PROC_LAMBDA_NAME = os.getenv("FILE_NAME_PROC_LAMBDA_NAME")
 
+def get_source_bucket_name() -> str:
+    """Get the SOURCE_BUCKET_NAME environment from environment variables."""
+    return os.getenv("SOURCE_BUCKET_NAME")
+
+def get_ack_bucket_name() -> str:
+    """Get the ACK_BUCKET_NAME environment from environment variables."""
+    return os.getenv("ACK_BUCKET_NAME")
 
 class FileStatus:
     """File status constants"""

ack_backend/src/constants.py

@@ -4,9 +4,9 @@
 from io import StringIO, BytesIO
 from typing import Union, Optional
 from botocore.exceptions import ClientError
-from constants import ACK_HEADERS, SOURCE_BUCKET_NAME, ACK_BUCKET_NAME, FILE_NAME_PROC_LAMBDA_NAME
+from constants import ACK_HEADERS, get_source_bucket_name, get_ack_bucket_name, FILE_NAME_PROC_LAMBDA_NAME
 from audit_table import change_audit_table_status_to_processed, get_next_queued_file_details
-from clients import s3_client, logger, lambda_client
+from clients import get_s3_client, logger, lambda_client
 from utils_for_ack_lambda import get_row_count
 from logging_decorators import upload_ack_file_logging_decorator
 
@@ -50,7 +50,7 @@ def obtain_current_ack_content(temp_ack_file_key: str) -> StringIO:
     """Returns the current ack file content if the file exists, or else initialises the content with the ack headers."""
     try:
         # If ack file exists in S3 download the contents
-        existing_ack_file = s3_client.get_object(Bucket=ACK_BUCKET_NAME, Key=temp_ack_file_key)
+        existing_ack_file = get_s3_client().get_object(Bucket=get_ack_bucket_name(), Key=temp_ack_file_key)
         existing_content = existing_ack_file["Body"].read().decode("utf-8")
     except ClientError as error:
         # If ack file does not exist in S3 create a new file containing the headers only
@@ -83,14 +83,19 @@ def upload_ack_file(
         cleaned_row = "|".join(data_row_str).replace(" |", "|").replace("| ", "|").strip()
         accumulated_csv_content.write(cleaned_row + "\n")
     csv_file_like_object = BytesIO(accumulated_csv_content.getvalue().encode("utf-8"))
-    s3_client.upload_fileobj(csv_file_like_object, ACK_BUCKET_NAME, temp_ack_file_key)
 
-    row_count_source = get_row_count(SOURCE_BUCKET_NAME, f"processing/{file_key}")
-    row_count_destination = get_row_count(ACK_BUCKET_NAME, temp_ack_file_key)
+    ack_bucket_name = get_ack_bucket_name()
+    source_bucket_name = get_source_bucket_name()
+
+    get_s3_client().upload_fileobj(csv_file_like_object, ack_bucket_name, temp_ack_file_key)
+
+    row_count_source = get_row_count(source_bucket_name, f"processing/{file_key}")
+    row_count_destination = get_row_count(ack_bucket_name, temp_ack_file_key)
     # TODO: Should we check for > and if so what handling is required
     if row_count_destination == row_count_source:
-        move_file(ACK_BUCKET_NAME, temp_ack_file_key, archive_ack_file_key)
-        move_file(SOURCE_BUCKET_NAME, f"processing/{file_key}", f"archive/{file_key}")
+        move_file(ack_bucket_name, temp_ack_file_key, archive_ack_file_key)
+        move_file(source_bucket_name, f"processing/{file_key}", f"archive/{file_key}")
+
         # Update the audit table and invoke the filename lambda with next file in the queue (if one exists)
         change_audit_table_status_to_processed(file_key, message_id)
         supplier_queue = f"{supplier}_{vaccine_type}"
@@ -107,7 +112,7 @@ def upload_ack_file(
         }
     else:
         result = None
-    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, archive_ack_file_key)
     return result
 
 
@@ -138,6 +143,7 @@ def update_ack_file(
 
 def move_file(bucket_name: str, source_file_key: str, destination_file_key: str) -> None:
     """Moves a file from one location to another within a single S3 bucket by copying and then deleting the file."""
+    s3_client = get_s3_client()
     s3_client.copy_object(
         Bucket=bucket_name, CopySource={"Bucket": bucket_name, "Key": source_file_key}, Key=destination_file_key
     )
@@ -150,7 +156,15 @@ def invoke_filename_lambda(file_key: str, message_id: str) -> None:
     try:
         lambda_payload = {
             "Records": [
-                {"s3": {"bucket": {"name": SOURCE_BUCKET_NAME}, "object": {"key": file_key}}, "message_id": message_id}
+                {"s3": 
+                    {
+                        "bucket": {
+                            "name": get_source_bucket_name()
+                        },
+                        "object": {"key": file_key}
+                    },
+                    "message_id": message_id
+                }
             ]
         }
         lambda_client.invoke(

ack_backend/src/update_ack_file.py

@@ -1,12 +1,12 @@
 """Utils for ack lambda"""
 
-from clients import s3_client
+from clients import get_s3_client
 
 
 def get_row_count(bucket_name: str, file_key: str) -> int:
     """
     Looks in the given bucket and returns the count of the number of lines in the given file.
     NOTE: Blank lines are not included in the count.
     """
-    response = s3_client.get_object(Bucket=bucket_name, Key=file_key)
+    response = get_s3_client().get_object(Bucket=bucket_name, Key=file_key)
     return sum(1 for line in response["Body"].iter_lines() if line.strip())

ack_backend/src/utils_for_ack_lambda.py

@@ -24,9 +24,6 @@
 with patch.dict("os.environ", MOCK_ENVIRONMENT_DICT):
     from ack_processor import lambda_handler
 
-s3_client = boto3_client("s3", region_name=REGION_NAME)
-firehose_client = boto3_client("firehose", region_name=REGION_NAME)
-
 BASE_SUCCESS_MESSAGE = MOCK_MESSAGE_DETAILS.success_message
 BASE_FAILURE_MESSAGE = {
     **{k: v for k, v in BASE_SUCCESS_MESSAGE.items() if k != "imms_id"},
@@ -41,19 +38,24 @@ class TestAckProcessor(unittest.TestCase):
     """Tests for the ack processor lambda handler."""
 
     def setUp(self) -> None:
-        GenericSetUp(s3_client, firehose_client)
+        self.s3_client = boto3_client("s3", region_name=REGION_NAME)
+        self.firehose_client = boto3_client("firehose", region_name=REGION_NAME)
+        GenericSetUp(self.s3_client, self.firehose_client)
 
         # MOCK SOURCE FILE WITH 100 ROWS TO SIMULATE THE SCENARIO WHERE THE ACK FILE IS NO FULL.
         # TODO: Test all other scenarios.
         mock_source_file_with_100_rows = StringIO("\n".join(f"Row {i}" for i in range(1, 101)))
-        s3_client.put_object(
+        self.s3_client.put_object(
             Bucket=BucketNames.SOURCE,
             Key=f"processing/{MOCK_MESSAGE_DETAILS.file_key}",
             Body=mock_source_file_with_100_rows.getvalue(),
         )
+        self.logger_info_patcher = patch('logging_decorators.logger.info')
+        self.mock_logger_info = self.logger_info_patcher.start()
 
     def tearDown(self) -> None:
-        GenericTearDown(s3_client, firehose_client)
+        GenericTearDown(self.s3_client, self.firehose_client)
+        self.mock_logger_info.stop()
 
     @staticmethod
     def generate_event(test_messages: list[dict]) -> dict:
@@ -111,7 +113,7 @@ def test_lambda_handler_main_multiple_records(self):
         response = lambda_handler(event=event, context={})
 
         self.assertEqual(response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
-        validate_ack_file_content(
+        validate_ack_file_content(self.s3_client,
             [*array_of_success_messages, *array_of_failure_messages, *array_of_mixed_success_and_failure_messages],
             existing_file_content=ValidValues.ack_headers,
         )
@@ -159,20 +161,20 @@ def test_lambda_handler_main(self):
             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)
-                validate_ack_file_content(test_case["messages"])
+                validate_ack_file_content(self.s3_client, test_case["messages"])
 
-                s3_client.delete_object(Bucket=BucketNames.DESTINATION, Key=MOCK_MESSAGE_DETAILS.temp_ack_file_key)
+                self.s3_client.delete_object(Bucket=BucketNames.DESTINATION, Key=MOCK_MESSAGE_DETAILS.temp_ack_file_key)
 
             # Test scenario where there is an existing ack file
             # TODO: None of the test cases have any existing ack file content?
             with self.subTest(msg=f"Existing ack file: {test_case['description']}"):
                 existing_ack_file_content = test_case.get("existing_ack_file_content", "")
-                setup_existing_ack_file(MOCK_MESSAGE_DETAILS.temp_ack_file_key, existing_ack_file_content)
+                setup_existing_ack_file(MOCK_MESSAGE_DETAILS.temp_ack_file_key, existing_ack_file_content, self.s3_client)
                 response = lambda_handler(event=self.generate_event(test_case["messages"]), context={})
                 self.assertEqual(response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
-                validate_ack_file_content(test_case["messages"], existing_ack_file_content)
+                validate_ack_file_content(self.s3_client, test_case["messages"], existing_ack_file_content)
 
-                s3_client.delete_object(Bucket=BucketNames.DESTINATION, Key=MOCK_MESSAGE_DETAILS.temp_ack_file_key)
+                self.s3_client.delete_object(Bucket=BucketNames.DESTINATION, Key=MOCK_MESSAGE_DETAILS.temp_ack_file_key)
 
     def test_lambda_handler_error_scenarios(self):
         """Test that the lambda handler raises appropriate exceptions for malformed event data."""

ack_backend/tests/test_ack_processor.py

@@ -0,0 +1,55 @@
+import unittest
+from unittest.mock import patch, MagicMock
+import audit_table
+from errors import UnhandledAuditTableError
+
+class TestAuditTable(unittest.TestCase):
+
+    def setUp(self):
+        self.logger_patcher = patch('audit_table.logger')
+        self.mock_logger = self.logger_patcher.start()
+        self.dynamodb_resource_patcher = patch('audit_table.dynamodb_resource')
+        self.mock_dynamodb_resource = self.dynamodb_resource_patcher.start()
+        self.dynamodb_client_patcher = patch('audit_table.dynamodb_client')
+        self.mock_dynamodb_client = self.dynamodb_client_patcher.start()
+
+    def tearDown(self):
+        self.logger_patcher.stop()
+        self.dynamodb_resource_patcher.stop()
+        self.dynamodb_client_patcher.stop()
+
+    def test_get_next_queued_file_details_returns_oldest(self):
+        # Arrange
+        mock_table = MagicMock()
+        self.mock_dynamodb_resource.Table.return_value = mock_table
+        mock_table.query.return_value = {
+            "Items": [
+                {"timestamp": 2, "my-key": "value2"},
+                {"timestamp": 1, "my-key": "value1"},
+            ]
+        }
+        # Act
+        result = audit_table.get_next_queued_file_details("queue1")
+        # Assert
+        self.assertEqual(result, {"timestamp": 1, "my-key": "value1"})
+
+    def test_get_next_queued_file_details_returns_none_if_empty(self):
+        mock_table = MagicMock()
+        self.mock_dynamodb_resource.Table.return_value = mock_table
+        mock_table.query.return_value = {"Items": []}
+        result = audit_table.get_next_queued_file_details("queue1")
+        self.assertIsNone(result)
+
+    def test_change_audit_table_status_to_processed_success(self):
+        # Should not raise
+        self.mock_dynamodb_client.update_item.return_value = {}
+        audit_table.change_audit_table_status_to_processed("file1", "msg1")
+        self.mock_dynamodb_client.update_item.assert_called_once()
+        self.mock_logger.info.assert_called_once()
+
+    def test_change_audit_table_status_to_processed_raises(self):
+        self.mock_dynamodb_client.update_item.side_effect = Exception("fail!")
+        with self.assertRaises(UnhandledAuditTableError) as ctx:
+            audit_table.change_audit_table_status_to_processed("file1", "msg1")
+        self.assertIn("fail!", str(ctx.exception))
+        self.mock_logger.error.assert_called_once()

ack_backend/tests/test_audit_table.py

@@ -27,10 +27,12 @@ class TestAckProcessor(unittest.TestCase):
     """Tests for the ack processor lambda handler."""
 
     def setUp(self) -> None:
-        GenericSetUp(s3_client, firehose_client)
+        self.s3_client = boto3_client("s3", region_name=REGION_NAME)
+        self.firehose_client = boto3_client("firehose", region_name=REGION_NAME)
+        GenericSetUp(self.s3_client, self.firehose_client)
 
     def tearDown(self) -> None:
-        GenericTearDown(s3_client, firehose_client)
+        GenericTearDown(self.s3_client, self.firehose_client)
 
     def test_get_error_message_for_ack_file(self):
         """Test the get_error_message_for_ack_file function."""