Adjust tests

Author: dlzhry2nhs

lambdas/ack_backend/tests/test_ack_processor.py

@@ -117,49 +117,47 @@ def assert_audit_entry_counts_equal(self, message_id: str, expected_counts: dict
             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")
+        actual_counts = {
+            "record_count": audit_entry.get("record_count", {}).get("N"),
+            "records_succeeded": audit_entry.get("records_succeeded", {}).get("N"),
+            "records_failed": audit_entry.get("records_failed", {}).get("N"),
+        }
 
-        self.assertEqual(actual_counts, expected_counts)
+        self.assertDictEqual(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
         add_audit_entry_to_table(self.dynamodb_client, "row")
-        # First array of messages: all successful. Rows 1 to 3
-        array_of_success_messages = [
+        # First array of messages. Rows 1 to 3
+        array_of_messages_one = [
             {
-                **BASE_SUCCESS_MESSAGE,
+                **BASE_FAILURE_MESSAGE,
                 "row_id": f"row^{i}",
-                "imms_id": f"imms_{i}",
                 "local_id": f"local^{i}",
             }
             for i in range(1, 4)
         ]
-        # Second array of messages: all with diagnostics (failure messages). Rows 4 to 7
-        array_of_failure_messages = [
+        # Second batch array of messages. Rows 4 to 7
+        array_of_messages_two = [
             {**BASE_FAILURE_MESSAGE, "row_id": f"row^{i}", "local_id": f"local^{i}"} for i in range(4, 8)
         ]
-        # Third array of messages: mixture of success and failure messages. Rows 8 to 11
-        array_of_mixed_success_and_failure_messages = [
+        # Third array of messages: mixture of diagnostic info.
+        array_of_messages_three = [
             {
                 **BASE_FAILURE_MESSAGE,
                 "row_id": "row^8",
                 "local_id": "local^8",
                 "diagnostics": DiagnosticsDictionaries.CUSTOM_VALIDATION_ERROR,
             },
             {
-                **BASE_SUCCESS_MESSAGE,
+                **BASE_FAILURE_MESSAGE,
                 "row_id": "row^9",
-                "imms_id": "imms_9",
                 "local_id": "local^9",
             },
             {
-                **BASE_SUCCESS_MESSAGE,
+                **BASE_FAILURE_MESSAGE,
                 "row_id": "row^10",
-                "imms_id": "imms_10",
                 "local_id": "local^10",
             },
             {
@@ -172,15 +170,15 @@ def test_lambda_handler_main_multiple_records(self):
 
         event = {
             "Records": [
-                {"body": json.dumps(array_of_success_messages)},
-                {"body": json.dumps(array_of_failure_messages)},
-                {"body": json.dumps(array_of_mixed_success_and_failure_messages)},
+                {"body": json.dumps(array_of_messages_one)},
+                {"body": json.dumps(array_of_messages_two)},
+                {"body": json.dumps(array_of_messages_three)},
             ]
         }
         expected_entry_counts = {
             "record_count": None,
             "records_succeeded": None,
-            "records_failed": "6",
+            "records_failed": "11",
         }
 
         response = lambda_handler(event=event, context={})
@@ -189,67 +187,58 @@ def test_lambda_handler_main_multiple_records(self):
         validate_ack_file_content(
             self.s3_client,
             [
-                *array_of_success_messages,
-                *array_of_failure_messages,
-                *array_of_mixed_success_and_failure_messages,
+                *array_of_messages_one,
+                *array_of_messages_two,
+                *array_of_messages_three,
             ],
             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."""
+        """Test lambda handler with consistent ack_file_name and message_template."""
         # Set up an audit entry which does not yet have record_count recorded
         add_audit_entry_to_table(self.dynamodb_client, "row")
         test_cases = [
-            {
-                "description": "Multiple messages: all successful",
-                "messages": [{"row_id": f"row^{i + 1}"} for i in range(10)],
-            },
             {
                 "description": "Multiple messages: all with diagnostics (failure messages)",
                 "messages": [
                     {"row_id": "row^1", "diagnostics": DiagnosticsDictionaries.UNIQUE_ID_MISSING},
                     {"row_id": "row^2", "diagnostics": DiagnosticsDictionaries.NO_PERMISSIONS},
                     {"row_id": "row^3", "diagnostics": DiagnosticsDictionaries.RESOURCE_NOT_FOUND_ERROR},
                 ],
+                "expected_failures_cum_tot": "3",
             },
             {
-                "description": "Multiple messages: mixture of success and failure messages",
+                "description": "Multiple messages: mixture of diagnostic outputs",
                 "messages": [
-                    {"row_id": "row^1", "imms_id": "TEST_IMMS_ID"},
-                    {"row_id": "row^2", "diagnostics": DiagnosticsDictionaries.UNIQUE_ID_MISSING},
+                    {"row_id": "row^1", "diagnostics": DiagnosticsDictionaries.UNIQUE_ID_MISSING},
+                    {"row_id": "row^2", "diagnostics": DiagnosticsDictionaries.CUSTOM_VALIDATION_ERROR},
                     {"row_id": "row^3", "diagnostics": DiagnosticsDictionaries.CUSTOM_VALIDATION_ERROR},
-                    {"row_id": "row^4"},
-                    {"row_id": "row^5", "diagnostics": DiagnosticsDictionaries.CUSTOM_VALIDATION_ERROR},
-                    {"row_id": "row^6", "diagnostics": DiagnosticsDictionaries.CUSTOM_VALIDATION_ERROR},
-                    {"row_id": "row^7"},
-                    {"row_id": "row^8", "diagnostics": DiagnosticsDictionaries.IDENTIFIER_DUPLICATION_ERROR},
+                    {"row_id": "row^4", "diagnostics": DiagnosticsDictionaries.CUSTOM_VALIDATION_ERROR},
+                    {"row_id": "row^5", "diagnostics": DiagnosticsDictionaries.IDENTIFIER_DUPLICATION_ERROR},
                 ],
-            },
-            {
-                "description": "Single row: success",
-                "messages": [{"row_id": "row^1"}],
+                "expected_failures_cum_tot": "8",
             },
             {
                 "description": "Single row: malformed diagnostics info from forwarder",
                 "messages": [{"row_id": "row^1", "diagnostics": "SHOULD BE A DICTIONARY, NOT A STRING"}],
+                "expected_failures_cum_tot": "9",
             },
         ]
-        expected_records_failed = [None, "3", "8", "8", "9"]
-        expected_entry_counts = {
-            "record_count": None,
-            "records_succeeded": None,
-            "records_failed": None,
-        }
 
-        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
+        for test_case in test_cases:
             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)
+                self.assert_audit_entry_counts_equal(
+                    "row",
+                    {
+                        "record_count": None,
+                        "records_succeeded": None,
+                        "records_failed": test_case["expected_failures_cum_tot"],
+                    },
+                )
                 validate_ack_file_content(self.s3_client, test_case["messages"])
 
                 self.s3_client.delete_object(
@@ -269,28 +258,27 @@ def test_lambda_handler_updates_ack_file_but_does_not_mark_complete_when_records
 
         # Original source file had 100 records
         add_audit_entry_to_table(self.dynamodb_client, mock_batch_message_id, record_count=100)
-        array_of_success_messages = [
+        array_of_failure_messages = [
             {
-                **BASE_SUCCESS_MESSAGE,
+                **BASE_FAILURE_MESSAGE,
                 "row_id": f"{mock_batch_message_id}^{i}",
-                "imms_id": f"imms_{i}",
                 "local_id": f"local^{i}",
             }
             for i in range(1, 4)
         ]
-        test_event = {"Records": [{"body": json.dumps(array_of_success_messages)}]}
+        test_event = {"Records": [{"body": json.dumps(array_of_failure_messages)}]}
         expected_entry_counts = {
             "record_count": "100",
             "records_succeeded": None,
-            "records_failed": None,
-        }
+            "records_failed": "3",
+        }  # Records succeeded not updated until all records are processed
 
         response = lambda_handler(event=test_event, context={})
 
         self.assertEqual(response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
         validate_ack_file_content(
             self.s3_client,
-            [*array_of_success_messages],
+            [*array_of_failure_messages],
             existing_file_content=ValidValues.ack_headers,
         )
         self.assert_ack_and_source_file_locations_correct(
@@ -324,31 +312,30 @@ def test_lambda_handler_updates_ack_file_and_marks_complete_when_all_records_pro
             Body=StringIO(existing_ack_content).getvalue(),
         )
 
-        array_of_success_messages = [
+        array_of_failure_messages = [
             {
-                **BASE_SUCCESS_MESSAGE,
+                **BASE_FAILURE_MESSAGE,
                 "row_id": f"{mock_batch_message_id}^{i}",
-                "imms_id": f"imms_{i}",
                 "local_id": f"local^{i}",
             }
             for i in range(50, 101)
         ]
 
         # Include the EoF message in the event
-        all_messages_plus_eof = deepcopy(array_of_success_messages)
+        all_messages_plus_eof = deepcopy(array_of_failure_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,
+            "records_succeeded": "49",
+            "records_failed": "51",
         }
 
         response = lambda_handler(event=test_event, context={})
 
         self.assertEqual(response, EXPECTED_ACK_LAMBDA_RESPONSE_FOR_SUCCESS)
         validate_ack_file_content(
-            self.s3_client, [*array_of_success_messages], existing_file_content=existing_ack_content, is_complete=True
+            self.s3_client, [*array_of_failure_messages], existing_file_content=existing_ack_content, is_complete=True
         )
         self.assert_ack_and_source_file_locations_correct(
             MOCK_MESSAGE_DETAILS.file_key,
@@ -359,60 +346,6 @@ def test_lambda_handler_updates_ack_file_and_marks_complete_when_all_records_pro
         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."""