Use contexts to open files

nationalarchives · Dec 19, 2024 · 7886c78 · 7886c78
1 parent 1db1e22
commit 7886c78
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 55 deletions.
diff --git 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
@@ -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 == "<error>This is the parser error log.</error>"
+        ) as tar:
+            result = lambda_function.create_parser_log_xml(tar)
+            assert result == "<error>This is the parser error log.</error>"
 
     @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 == "<error>parser.log not found</error>"
+        ) as tar:
+            tar.extractfile = MagicMock(side_effect=KeyError)
+            result = lambda_function.create_parser_log_xml(tar)
+            assert result == "<error>parser.log not found</error>"
 
     @patch.dict(
         os.environ,