DAS-NONE: Use TemporaryDirectory context in process_item

nasa · Dec 12, 2024 · 4fcded7 · 4fcded7
1 parent 54049a8
commit 4fcded7
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 71 deletions.
diff --git a/harmony_service/adapter.py b/harmony_service/adapter.py
@@ -8,8 +8,7 @@
 """
 
 from pathlib import Path
-from shutil import rmtree
-from tempfile import mkdtemp
+from tempfile import TemporaryDirectory
 
 from harmony_service_lib import BaseHarmonyAdapter
 from harmony_service_lib.message import Source as HarmonySource
@@ -93,61 +92,59 @@ def get_asset_from_item(self, item: Item) -> Asset:
 
     def process_item(self, item: Item, source: HarmonySource) -> Item:
         """Processes a single input STAC item."""
-        try:
-            working_directory = mkdtemp()
-            results = item.clone()
-            results.assets = {}
-
-            asset = self.get_asset_from_item(item)
-
-            color_palette = get_color_palette_from_item(item)
-
-            # Download the input:
-            input_data_filename = download(
-                asset.href,
-                working_directory,
-                logger=self.logger,
-                cfg=self.config,
-                access_token=self.message.accessToken,
-            )
+        with TemporaryDirectory() as working_directory:
+            try:
+                results = item.clone()
+                results.assets = {}
+
+                asset = self.get_asset_from_item(item)
+
+                color_palette = get_color_palette_from_item(item)
+
+                # Download the input:
+                input_data_filename = download(
+                    asset.href,
+                    working_directory,
+                    logger=self.logger,
+                    cfg=self.config,
+                    access_token=self.message.accessToken,
+                )
 
-            # Create browse images.
-            image_file_list = create_browse_imagery(
-                self.message,
-                input_data_filename,
-                source,
-                color_palette,
-                logger=self.logger,
-            )
+                # Create browse images.
+                image_file_list = create_browse_imagery(
+                    self.message,
+                    input_data_filename,
+                    source,
+                    color_palette,
+                    logger=self.logger,
+                )
 
-            # image_file_list is a list of tuples (image, world, auxiliary)
-            # we need to stage them each individually, and then add their final
-            # locations to a list before creating the stac item.
-            item_assets = []
-
-            for (
-                browse_image_name,
-                world_file_name,
-                aux_xml_file_name,
-            ) in image_file_list:
-                # Stage the images:
-                browse_image_url = self.stage_output(browse_image_name, asset.href)
-                browse_aux_url = self.stage_output(aux_xml_file_name, asset.href)
-                world_file_url = self.stage_output(world_file_name, asset.href)
-                item_assets.append(('data', browse_image_url, 'data'))
-                item_assets.append(('metadata', world_file_url, 'metadata'))
-                item_assets.append(('auxiliary', browse_aux_url, 'metadata'))
-
-            manifest_url = self.stage_manifest(image_file_list, asset.href)
-            item_assets.insert(0, ('data', manifest_url, 'metadata'))
-
-            return self.create_output_stac_item(item, item_assets)
-
-        except Exception as exception:
-            self.logger.exception(exception)
-            raise HyBIGServiceError(str(exception)) from exception
-        finally:
-            rmtree(working_directory)
+                # image_file_list is a list of tuples (image, world, auxiliary)
+                # we need to stage them each individually, and then add their final
+                # locations to a list before creating the stac item.
+                item_assets = []
+
+                for (
+                    browse_image_name,
+                    world_file_name,
+                    aux_xml_file_name,
+                ) in image_file_list:
+                    # Stage the images:
+                    browse_image_url = self.stage_output(browse_image_name, asset.href)
+                    browse_aux_url = self.stage_output(aux_xml_file_name, asset.href)
+                    world_file_url = self.stage_output(world_file_name, asset.href)
+                    item_assets.append(('data', browse_image_url, 'data'))
+                    item_assets.append(('metadata', world_file_url, 'metadata'))
+                    item_assets.append(('auxiliary', browse_aux_url, 'metadata'))
+
+                manifest_url = self.stage_manifest(image_file_list, asset.href)
+                item_assets.insert(0, ('data', manifest_url, 'metadata'))
+
+                return self.create_output_stac_item(item, item_assets)
+
+            except Exception as exception:
+                self.logger.exception(exception)
+                raise HyBIGServiceError(str(exception)) from exception
 
     def stage_output(self, transformed_file: Path, input_file: str) -> str:
         """Generate an output file name based on the input asset URL and the

diff --git a/tests/test_service/test_adapter.py b/tests/test_service/test_adapter.py
@@ -100,18 +100,19 @@ def assert_expected_output_catalog(
             },
         )
 
+    @patch('harmony_service.adapter.TemporaryDirectory')
     @patch('hybig.browse.reproject')
-    @patch('harmony_service.adapter.rmtree')
-    @patch('harmony_service.adapter.mkdtemp')
     @patch('harmony_service.adapter.download')
     @patch('harmony_service.adapter.stage')
     def test_valid_request_jpeg(
-        self, mock_stage, mock_download, mock_mkdtemp, mock_rmtree, mock_reproject
+        self, mock_stage, mock_download, mock_reproject, mock_temp_dir
     ):
         """Ensure a request with a correctly formatted message is fully
         processed.
 
         """
+        mock_temp_dir.return_value.__enter__.return_value = self.temp_dir
+
         expected_downloaded_file = self.temp_dir / 'input.tiff'
 
         expected_browse_basename = 'input.jpg'
@@ -136,8 +137,6 @@ def test_valid_request_jpeg(
         expected_aux_mime = 'application/xml'
         expected_manifest_mime = 'text/plain'
 
-        mock_mkdtemp.return_value = self.temp_dir
-
         def move_tif(*args, **kwargs):
             """Copy fixture tiff to download location."""
             copy(self.red_tif_fixture, expected_downloaded_file)
@@ -332,21 +331,19 @@ def move_tif(*args, **kwargs):
             ]
         )
 
-        # Ensure container clean-up was requested:
-        mock_rmtree.assert_called_once_with(self.temp_dir)
-
+    @patch('harmony_service.adapter.TemporaryDirectory')
     @patch('hybig.browse.reproject')
-    @patch('harmony_service.adapter.rmtree')
-    @patch('harmony_service.adapter.mkdtemp')
     @patch('harmony_service.adapter.download')
     @patch('harmony_service.adapter.stage')
     def test_valid_request_png(
-        self, mock_stage, mock_download, mock_mkdtemp, mock_rmtree, mock_reproject
+        self, mock_stage, mock_download, mock_reproject, mock_temp_dir
     ):
         """Ensure a request with a correctly formatted message is fully
         processed.
 
         """
+        mock_temp_dir.return_value.__enter__.return_value = self.temp_dir
+
         expected_downloaded_file = self.temp_dir / 'input.tiff'
 
         expected_browse_basename = 'input.png'
@@ -371,8 +368,6 @@ def test_valid_request_png(
         expected_aux_mime = 'application/xml'
         expected_manifest_mime = 'text/plain'
 
-        mock_mkdtemp.return_value = self.temp_dir
-
         def move_tif(*args, **kwargs):
             """Copy fixture tiff to download location."""
             copy(self.red_tif_fixture, expected_downloaded_file)
@@ -552,9 +547,6 @@ def move_tif(*args, **kwargs):
             ]
         )
 
-        # Ensure container clean-up was requested:
-        mock_rmtree.assert_called_once_with(self.temp_dir)
-
     @patch('harmony_service.adapter.download')
     def test_forbidden_download(self, mock_download):
         mock_download.side_effect = ForbiddenException('You are forbidden to download.')