Skip to content

Commit e7bbcc6

Browse files
JamesW1-NHSJamesW1-NHS
committed
audit_table & unit tests
1 parent 041ea67 commit e7bbcc6

File tree

2 files changed

+106
-13
lines changed

2 files changed

+106
-13
lines changed

‎lambdas/ack_backend/src/audit_table.py‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,20 +11,21 @@ def change_audit_table_status_to_processed(file_key: str, message_id: str) -> No
1111
"""Updates the status in the audit table to 'Processed' and returns the queue name."""
1212
try:
1313
# Update the status in the audit table to "Processed"
14-
dynamodb_client.update_item(
14+
response = dynamodb_client.update_item(
1515
TableName=AUDIT_TABLE_NAME,
1616
Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}},
1717
UpdateExpression="SET #status = :status",
1818
ExpressionAttributeNames={"#status": "status"},
1919
ExpressionAttributeValues={":status": {"S": FileStatus.PROCESSED}},
2020
ConditionExpression="attribute_exists(message_id)",
21+
ReturnValues="UPDATED_NEW",
2122
)
22-
23+
result = response.get("Attributes", {}).get("status").get("S")
2324
logger.info(
2425
"The status of %s file, with message id %s, was successfully updated to %s in the audit table",
2526
file_key,
2627
message_id,
27-
FileStatus.PROCESSED,
28+
result,
2829
)
2930

3031
except Exception as error: # pylint: disable = broad-exception-caught
@@ -46,7 +47,7 @@ def get_record_count_by_message_id(event_message_id: str) -> Optional[int]:
4647
return int(record_count)
4748

4849

49-
def set_record_success_count(message_id: str) -> Optional[int]:
50+
def set_records_succeeded_count(message_id: str) -> None:
5051
"""Set the 'records_succeeded' item in the audit table entry"""
5152
audit_record = dynamodb_client.get_item(
5253
TableName=AUDIT_TABLE_NAME, Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}}
@@ -58,52 +59,56 @@ def set_record_success_count(message_id: str) -> Optional[int]:
5859
records_failed = int(records_failed_item) if records_failed_item else 0
5960
records_succeeded = record_count - records_failed
6061

62+
counter_name = AuditTableKeys.RECORDS_SUCCEEDED
6163
try:
6264
response = dynamodb_client.update_item(
6365
TableName=AUDIT_TABLE_NAME,
6466
Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}},
6567
UpdateExpression="SET #counter = :value",
66-
ExpressionAttributeNames={"#counter": AuditTableKeys.RECORDS_SUCCEEDED},
67-
ExpressionAttributeValues={":value": {"S": str(records_succeeded)}},
68+
ExpressionAttributeNames={"#counter": counter_name},
69+
ExpressionAttributeValues={":value": {"N": str(records_succeeded)}},
6870
ConditionExpression="attribute_exists(message_id)",
6971
ReturnValues="UPDATED_NEW",
7072
)
73+
result = response.get("Attributes", {}).get(counter_name).get("N")
7174
logger.info(
7275
"Counter %s for message id %s set to %s in the audit table",
73-
AuditTableKeys.RECORDS_SUCCEEDED,
76+
counter_name,
7477
message_id,
75-
response.get("Attributes", {}).get(AuditTableKeys.RECORDS_SUCCEEDED),
78+
result,
7679
)
7780

7881
except Exception as error: # pylint: disable = broad-exception-caught
7982
logger.error(error)
8083
raise UnhandledAuditTableError(error) from error
8184

8285

83-
def increment_record_counter(message_id: str, counter_name: str = AuditTableKeys.RECORDS_FAILED) -> None:
86+
def increment_records_failed_count(message_id: str) -> None:
8487
"""
8588
Increment a counter attribute safely, handling the case where it might not exist.
8689
From https://docs.aws.amazon.com/code-library/latest/ug/dynamodb_example_dynamodb_Scenario_AtomicCounterOperations_section.html
8790
"""
8891

8992
increment_value = 1
9093
initial_value = 0
94+
counter_name = AuditTableKeys.RECORDS_FAILED
9195
try:
9296
# Use SET with if_not_exists to safely increment the counter
9397
response = dynamodb_client.update_item(
9498
TableName=AUDIT_TABLE_NAME,
9599
Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}},
96100
UpdateExpression="SET #counter = if_not_exists(#counter, :initial) + :increment",
97101
ExpressionAttributeNames={"#counter": counter_name},
98-
ExpressionAttributeValues={":increment": {"S": str(increment_value)}, ":initial": {"S": str(initial_value)}},
102+
ExpressionAttributeValues={":increment": {"N": str(increment_value)}, ":initial": {"N": str(initial_value)}},
99103
ConditionExpression="attribute_exists(message_id)",
100104
ReturnValues="UPDATED_NEW",
101105
)
106+
result = response.get("Attributes", {}).get(counter_name).get("N")
102107
logger.info(
103108
"Counter %s for message id %s incremented to %s in the audit table",
104109
counter_name,
105110
message_id,
106-
response.get("Attributes", {}).get(counter_name),
111+
result,
107112
)
108113

109114
except Exception as error: # pylint: disable = broad-exception-caught

‎lambdas/ack_backend/tests/test_audit_table.py‎

Lines changed: 90 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import audit_table
55
from common.models.errors import UnhandledAuditTableError
6+
from constants import AUDIT_TABLE_NAME, AuditTableKeys, FileStatus
67

78

89
class TestAuditTable(unittest.TestCase):
@@ -17,9 +18,16 @@ def tearDown(self):
1718

1819
def test_change_audit_table_status_to_processed_success(self):
1920
# Should not raise
20-
self.mock_dynamodb_client.update_item.return_value = {}
2121
audit_table.change_audit_table_status_to_processed("file1", "msg1")
22-
self.mock_dynamodb_client.update_item.assert_called_once()
22+
self.mock_dynamodb_client.update_item.assert_called_once_with(
23+
TableName=AUDIT_TABLE_NAME,
24+
Key={AuditTableKeys.MESSAGE_ID: {"S": "msg1"}},
25+
UpdateExpression="SET #status = :status",
26+
ExpressionAttributeNames={"#status": "status"},
27+
ExpressionAttributeValues={":status": {"S": FileStatus.PROCESSED}},
28+
ConditionExpression="attribute_exists(message_id)",
29+
ReturnValues="UPDATED_NEW",
30+
)
2331
self.mock_logger.info.assert_called_once()
2432

2533
def test_change_audit_table_status_to_processed_raises(self):
@@ -46,3 +54,83 @@ def test_get_record_count_by_message_id_returns_none_if_record_count_not_set(sel
4654
self.mock_dynamodb_client.get_item.return_value = {"Item": {"message_id": {"S": test_message_id}}}
4755

4856
self.assertIsNone(audit_table.get_record_count_by_message_id(test_message_id))
57+
58+
def test_set_records_succeeded_count(self):
59+
test_message_id = "1234"
60+
self.mock_dynamodb_client.get_item.return_value = {
61+
"Item": {"message_id": {"S": test_message_id}, "record_count": {"N": "1000"}, "records_failed": {"N": "42"}}
62+
}
63+
audit_table.set_records_succeeded_count(test_message_id)
64+
self.mock_dynamodb_client.get_item.assert_called_once()
65+
self.mock_dynamodb_client.update_item.assert_called_once_with(
66+
TableName=AUDIT_TABLE_NAME,
67+
Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
68+
UpdateExpression="SET #counter = :value",
69+
ExpressionAttributeNames={"#counter": AuditTableKeys.RECORDS_SUCCEEDED},
70+
ExpressionAttributeValues={":value": {"N": "958"}},
71+
ConditionExpression="attribute_exists(message_id)",
72+
ReturnValues="UPDATED_NEW",
73+
)
74+
self.mock_logger.info.assert_called_once()
75+
76+
def test_set_records_succeeded_count_no_failures(self):
77+
test_message_id = "1234"
78+
self.mock_dynamodb_client.get_item.return_value = {
79+
"Item": {"message_id": {"S": test_message_id}, "record_count": {"N": "1000"}}
80+
}
81+
audit_table.set_records_succeeded_count(test_message_id)
82+
self.mock_dynamodb_client.get_item.assert_called_once()
83+
self.mock_dynamodb_client.update_item.assert_called_once_with(
84+
TableName=AUDIT_TABLE_NAME,
85+
Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
86+
UpdateExpression="SET #counter = :value",
87+
ExpressionAttributeNames={"#counter": AuditTableKeys.RECORDS_SUCCEEDED},
88+
ExpressionAttributeValues={":value": {"N": "1000"}},
89+
ConditionExpression="attribute_exists(message_id)",
90+
ReturnValues="UPDATED_NEW",
91+
)
92+
self.mock_logger.info.assert_called_once()
93+
94+
def test_set_records_succeeded_count_no_records(self):
95+
test_message_id = "1234"
96+
self.mock_dynamodb_client.get_item.return_value = {"Item": {"message_id": {"S": test_message_id}}}
97+
audit_table.set_records_succeeded_count(test_message_id)
98+
self.mock_dynamodb_client.get_item.assert_called_once()
99+
self.mock_dynamodb_client.update_item.assert_called_once_with(
100+
TableName=AUDIT_TABLE_NAME,
101+
Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
102+
UpdateExpression="SET #counter = :value",
103+
ExpressionAttributeNames={"#counter": AuditTableKeys.RECORDS_SUCCEEDED},
104+
ExpressionAttributeValues={":value": {"N": "0"}},
105+
ConditionExpression="attribute_exists(message_id)",
106+
ReturnValues="UPDATED_NEW",
107+
)
108+
self.mock_logger.info.assert_called_once()
109+
110+
def test_set_records_succeeded_count_raises(self):
111+
self.mock_dynamodb_client.update_item.side_effect = Exception("fail!")
112+
with self.assertRaises(UnhandledAuditTableError) as ctx:
113+
audit_table.set_records_succeeded_count("msg1")
114+
self.assertIn("fail!", str(ctx.exception))
115+
self.mock_logger.error.assert_called_once()
116+
117+
def test_increment_records_failed_count(self):
118+
test_message_id = "1234"
119+
audit_table.increment_records_failed_count(test_message_id)
120+
self.mock_dynamodb_client.update_item.assert_called_once_with(
121+
TableName=AUDIT_TABLE_NAME,
122+
Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
123+
UpdateExpression="SET #counter = if_not_exists(#counter, :initial) + :increment",
124+
ExpressionAttributeNames={"#counter": AuditTableKeys.RECORDS_FAILED},
125+
ExpressionAttributeValues={":increment": {"N": "1"}, ":initial": {"N": "0"}},
126+
ConditionExpression="attribute_exists(message_id)",
127+
ReturnValues="UPDATED_NEW",
128+
)
129+
self.mock_logger.info.assert_called_once()
130+
131+
def test_increment_records_failed_count_raises(self):
132+
self.mock_dynamodb_client.update_item.side_effect = Exception("fail!")
133+
with self.assertRaises(UnhandledAuditTableError) as ctx:
134+
audit_table.increment_records_failed_count("msg1")
135+
self.assertIn("fail!", str(ctx.exception))
136+
self.mock_logger.error.assert_called_once()

0 commit comments

Comments
 (0)