From 7886c78655d91877a016e7e14da10ceb2556153e Mon Sep 17 00:00:00 2001 From: Nick Jackson Date: Thu, 19 Dec 2024 12:41:37 +0000 Subject: [PATCH] Use contexts to open files --- ds-caselaw-ingester/lambda_function.py | 16 +++-- ds-caselaw-ingester/tests.py | 96 +++++++++++++------------- 2 files changed, 57 insertions(+), 55 deletions(-) diff --git a/ds-caselaw-ingester/lambda_function.py b/ds-caselaw-ingester/lambda_function.py index 7b85789..c7fc1bb 100644 --- a/ds-caselaw-ingester/lambda_function.py +++ b/ds-caselaw-ingester/lambda_function.py @@ -392,7 +392,8 @@ def __init__(self, message: Message): print(f"Ingester Start: Consignment reference {self.consignment_reference}") print(f"Received Message: {self.message.message}") self.local_tar_filename = self.save_tar_file_in_s3() - self.tar = tarfile.open(self.local_tar_filename, mode="r") + with tarfile.open(self.local_tar_filename, mode="r") as tar: + self.tar = tar self.metadata = extract_metadata(self.tar, self.consignment_reference) self.message.update_consignment_reference(self.metadata["parameters"]["TRE"]["reference"]) self.consignment_reference = self.message.get_consignment_reference() @@ -532,12 +533,13 @@ def save_files_to_s3(self) -> None: modified_targz_filename = ( self.local_tar_filename if docx_filename else modify_filename(self.local_tar_filename, "_nodocx") ) - store_file( - open(self.local_tar_filename, mode="rb"), - self.uri, - os.path.basename(modified_targz_filename), - s3_client, - ) + with open(self.local_tar_filename, mode="rb") as local_tar: + store_file( + local_tar, + self.uri, + os.path.basename(modified_targz_filename), + s3_client, + ) print(f"saved tar.gz as {modified_targz_filename!r}") # Store docx and rename diff --git a/ds-caselaw-ingester/tests.py b/ds-caselaw-ingester/tests.py index 90fe221..d2a2a24 100644 --- a/ds-caselaw-ingester/tests.py +++ b/ds-caselaw-ingester/tests.py @@ -243,50 +243,50 @@ class TestLambda: ) def test_extract_xml_file_success_tdr(self): - filename = "TDR-2022-DNWR.xml" - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - result = lambda_function.extract_xml_file(tar, filename) - xml = ET.XML(result.read()) - assert xml.tag == "{http://docs.oasis-open.org/legaldocml/ns/akn/3.0}akomaNtoso" + ) as tar: + filename = "TDR-2022-DNWR.xml" + result = lambda_function.extract_xml_file(tar, filename) + xml = ET.XML(result.read()) + assert xml.tag == "{http://docs.oasis-open.org/legaldocml/ns/akn/3.0}akomaNtoso" def test_extract_xml_file_not_found_tdr(self): - filename = "unknown.xml" - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - result = lambda_function.extract_xml_file(tar, filename) - assert result is None + ) as tar: + filename = "unknown.xml" + result = lambda_function.extract_xml_file(tar, filename) + assert result is None def test_extract_xml_file_name_empty(self): - filename = "" - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - result = lambda_function.extract_xml_file(tar, filename) - assert result is None + ) as tar: + filename = "" + result = lambda_function.extract_xml_file(tar, filename) + assert result is None def test_extract_metadata_success_tdr(self): - consignment_reference = "TDR-2022-DNWR" - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - result = lambda_function.extract_metadata(tar, consignment_reference) - assert result["parameters"]["TRE"]["payload"] is not None + ) as tar: + consignment_reference = "TDR-2022-DNWR" + result = lambda_function.extract_metadata(tar, consignment_reference) + assert result["parameters"]["TRE"]["payload"] is not None def test_extract_metadata_not_found_tdr(self): - consignment_reference = "unknown_consignment_reference" - tar = tarfile.open( + with tarfile.open( self.TARBALL_MISSING_METADATA_PATH, mode="r", - ) - with pytest.raises(lambda_function.FileNotFoundException, match="Consignment Ref:"): - lambda_function.extract_metadata(tar, consignment_reference) + ) as tar: + consignment_reference = "unknown_consignment_reference" + with pytest.raises(lambda_function.FileNotFoundException, match="Consignment Ref:"): + lambda_function.extract_metadata(tar, consignment_reference) def test_extract_docx_filename_success(self): metadata = {"parameters": {"TRE": {"payload": {"filename": "judgment.docx"}}}} @@ -496,43 +496,43 @@ def test_send_updated_judgment_notification_with_no_tdr_section(self, mock_print @patch.object(lambda_function, "store_file") def test_copy_file_success(self, mock_store_file): - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - filename = "TDR-2022-DNWR/TDR-2022-DNWR.xml" - session = boto3.Session - lambda_function.store_file = MagicMock() - lambda_function.copy_file(tar, filename, "new_filename", "uri", session) - lambda_function.store_file.assert_called_with(ANY, ANY, ANY, ANY) + ) as tar: + filename = "TDR-2022-DNWR/TDR-2022-DNWR.xml" + session = boto3.Session + lambda_function.store_file = MagicMock() + lambda_function.copy_file(tar, filename, "new_filename", "uri", session) + lambda_function.store_file.assert_called_with(ANY, ANY, ANY, ANY) def test_copy_file_not_found(self): - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - filename = "does_not_exist.txt" - session = boto3.Session - with pytest.raises(lambda_function.FileNotFoundException): - lambda_function.copy_file(tar, filename, "new_filename", "uri", session) + ) as tar: + filename = "does_not_exist.txt" + session = boto3.Session + with pytest.raises(lambda_function.FileNotFoundException): + lambda_function.copy_file(tar, filename, "new_filename", "uri", session) def test_create_xml_contents_success(self): - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - result = lambda_function.create_parser_log_xml(tar) - assert result == "This is the parser error log." + ) as tar: + result = lambda_function.create_parser_log_xml(tar) + assert result == "This is the parser error log." @patch.object(tarfile, "open") def test_create_xml_contents_failure(self, mock_open_tarfile): - tar = tarfile.open( + with tarfile.open( self.TDR_TARBALL_PATH, mode="r", - ) - tar.extractfile = MagicMock(side_effect=KeyError) - result = lambda_function.create_parser_log_xml(tar) - assert result == "parser.log not found" + ) as tar: + tar.extractfile = MagicMock(side_effect=KeyError) + result = lambda_function.create_parser_log_xml(tar) + assert result == "parser.log not found" @patch.dict( os.environ,