[PRMP-990] update tests for stitched reference method to include S3 version id

Author: NogaNHS

lambdas/services/pdf_stitching_service.py

@@ -45,7 +45,7 @@ def __init__(self):
         self.document_service = DocumentService()
         self.sqs_service = SQSService()
         self.multipart_references: list[DocumentReference] = []
-        self.stitched_reference: DocumentReference = None
+        self.stitched_reference: DocumentReference | None = None
 
     def retrieve_multipart_references(
         self, nhs_number: str, doc_type: SupportedDocumentTypes
@@ -101,7 +101,7 @@ def process_message(self, stitching_message: PdfStitchingSqsMessage):
                 stitch_file_size=sys.getsizeof(stitching_data_stream),
             )
             self.upload_stitched_file(stitching_data_stream=stitching_data_stream)
-            self.update_stitched_reference()
+            self.update_stitched_reference_with_version_id()
             self.migrate_multipart_references()
             self.write_stitching_reference()
             self.publish_nrl_message(
@@ -130,10 +130,9 @@ def create_stitched_reference(
             deep=True,
         )
 
-    def update_stitched_reference(self):
+    def update_stitched_reference_with_version_id(self):
         self.stitched_reference.s3_version_id = self.s3_service.get_head_object(
-            self.target_bucket, 
-            self.stitched_reference.s3_file_key
+            self.target_bucket, self.stitched_reference.s3_file_key
         ).get("VersionId")
 
     def process_stitching(self, s3_object_keys: list[str]) -> BytesIO:
@@ -377,6 +376,6 @@ def process_manual_trigger(self, ods_code: str, queue_url):
             except (InvalidMessageException, Exception) as e:
                 logger.error(f"Error sending batch to SQS: {str(e)}")
             # 1 batch is 10 messages
-            # we can send up to 300 messages a second
-            # a 0.1s delay means 10 batches, so 100 messages
+            # we can send up to 300 messages a second.
+            # a 0.1 s delay means 10 batches, so 100 messages
             time.sleep(0.1)

lambdas/tests/unit/services/test_pdf_stitching_service.py

@@ -71,6 +71,13 @@ def mock_upload_stitched_file(mocker, mock_service):
     return mocker.patch.object(mock_service, "upload_stitched_file")
 
 
+@pytest.fixture
+def mock_update_stitched_reference_with_version_id(mocker, mock_service):
+    return mocker.patch.object(
+        mock_service, "update_stitched_reference_with_version_id"
+    )
+
+
 @pytest.fixture
 def mock_migrate_multipart_references(mocker, mock_service):
     return mocker.patch.object(mock_service, "migrate_multipart_references")
@@ -164,6 +171,7 @@ def test_process_message(
     mock_sort_multipart_object_keys,
     mock_process_stitching,
     mock_upload_stitched_file,
+    mock_update_stitched_reference_with_version_id,
     mock_migrate_multipart_references,
     mock_write_stitching_reference,
     mock_publish_nrl_message,
@@ -193,6 +201,7 @@ def set_stitched_reference(document_reference, stitch_file_size, *args, **kwargs
     mock_sort_multipart_object_keys.assert_called_once_with()
     mock_process_stitching.assert_called_once_with(s3_object_keys=test_sorted_keys)
     mock_upload_stitched_file.assert_called_once_with(stitching_data_stream=test_stream)
+    mock_update_stitched_reference_with_version_id.assert_called_once()
     mock_migrate_multipart_references.assert_called_once()
     mock_write_stitching_reference.assert_called_once()
     mock_publish_nrl_message.assert_called_once()
@@ -671,6 +680,24 @@ def test_rollback_reference_migration_handles_exception(mock_service):
         mock_service.rollback_reference_migration()
 
 
+def test_update_stitched_reference_with_version_id(mock_service):
+    test_version_id = "test-version-id-12345"
+    mock_service.stitched_reference = TEST_1_OF_1_DOCUMENT_REFERENCE
+
+    mock_service.s3_service.get_head_object.return_value = {
+        "VersionId": test_version_id,
+        "ContentType": "application/pdf",
+        "ContentLength": 1234,
+    }
+
+    mock_service.update_stitched_reference_with_version_id()
+
+    mock_service.s3_service.get_head_object.assert_called_once_with(
+        MOCK_LG_BUCKET, TEST_1_OF_1_DOCUMENT_REFERENCE.s3_file_key
+    )
+    assert mock_service.stitched_reference.s3_version_id == test_version_id
+
+
 def test_process_manual_trigger_calls_process_message_for_each_nhs_number(
     mocker, mock_service
 ):