Refactor and simplify assembly workflow. (#827)

* Extracted assemble args into a class, add keep.

Signed-off-by: dblock <[email protected]>

* Remove TemporaryDirectory inside another TemporaryDirectory.

Signed-off-by: dblock <[email protected]>

src/assemble_workflow/assemble_args.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+
+import argparse
+import logging
+
+
+class AssembleArgs:
+    manifest: str
+    keep: bool
+
+    def __init__(self):
+        parser = argparse.ArgumentParser(description="Assemble an OpenSearch Distribution")
+        parser.add_argument("manifest", type=argparse.FileType("r"), help="Manifest file.")
+        parser.add_argument("-b", "--base-url", dest="base_url", help="The base url to download the artifacts.")
+        parser.add_argument(
+            "--keep",
+            dest="keep",
+            action="store_true",
+            help="Do not delete the working temporary directory.",
+        )
+        parser.add_argument(
+            "-v",
+            "--verbose",
+            help="Show more verbose output.",
+            action="store_const",
+            default=logging.INFO,
+            const=logging.DEBUG,
+            dest="logging_level",
+        )
+
+        args = parser.parse_args()
+        self.logging_level = args.logging_level
+        self.manifest = args.manifest
+        self.keep = args.keep
+        self.base_url = args.base_url

src/assemble_workflow/bundle.py

@@ -29,7 +29,7 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_traceback):
         self.tmp_dir.__exit__(exc_type, exc_value, exc_traceback)
 
-    def __init__(self, build_manifest, artifacts_dir, bundle_recorder):
+    def __init__(self, build_manifest, artifacts_dir, bundle_recorder, keep=False):
         """
         Construct a new Bundle instance.
         :param build_manifest: A BuildManifest created from the build workflow.
@@ -39,7 +39,7 @@ def __init__(self, build_manifest, artifacts_dir, bundle_recorder):
         self.plugins = self.__get_plugins(build_manifest.components.values())
         self.artifacts_dir = artifacts_dir
         self.bundle_recorder = bundle_recorder
-        self.tmp_dir = TemporaryDirectory()
+        self.tmp_dir = TemporaryDirectory(keep=keep)
         self.min_dist = self.__get_min_dist(build_manifest.components.values())
         self.installed_plugins = []
 

src/assemble_workflow/bundles.py

@@ -22,6 +22,6 @@ def from_name(cls, name):
         return klass
 
     @classmethod
-    def create(cls, build_manifest, artifacts_dir, bundle_recorder):
+    def create(cls, build_manifest, artifacts_dir, bundle_recorder, keep):
         klass = cls.from_name(build_manifest.build.name)
-        return klass(build_manifest, artifacts_dir, bundle_recorder)
+        return klass(build_manifest, artifacts_dir, bundle_recorder, keep)

src/run_assemble.py

@@ -6,32 +6,19 @@
 # this file be licensed under the Apache-2.0 license or a
 # compatible open source license.
 
-import argparse
 import logging
 import os
 import sys
 
+from assemble_workflow.assemble_args import AssembleArgs
 from assemble_workflow.bundle_recorder import BundleRecorder
 from assemble_workflow.bundles import Bundles
 from manifests.build_manifest import BuildManifest
 from system import console
-from system.temporary_directory import TemporaryDirectory
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Assemble an OpenSearch Bundle")
-    parser.add_argument("manifest", type=argparse.FileType("r"), help="Manifest file.")
-    parser.add_argument(
-        "-v",
-        "--verbose",
-        help="Show more verbose output.",
-        action="store_const",
-        default=logging.INFO,
-        const=logging.DEBUG,
-        dest="logging_level",
-    )
-    parser.add_argument("-b", "--base-url", dest="base_url", help="The base url to download the artifacts.")
-    args = parser.parse_args()
+    args = AssembleArgs()
 
     console.configure(level=args.logging_level)
 
@@ -41,21 +28,20 @@ def main():
     output_dir = os.path.join(os.getcwd(), "dist")
     os.makedirs(output_dir, exist_ok=True)
 
-    with TemporaryDirectory(chdir=True):
-        logging.info(f"Bundling {build.name} ({build.architecture}) on {build.platform} into {output_dir} ...")
+    logging.info(f"Bundling {build.name} ({build.architecture}) on {build.platform} into {output_dir} ...")
 
-        bundle_recorder = BundleRecorder(build, output_dir, artifacts_dir, args.base_url)
+    bundle_recorder = BundleRecorder(build, output_dir, artifacts_dir, args.base_url)
 
-        with Bundles.create(build_manifest, artifacts_dir, bundle_recorder) as bundle:
-            bundle.install_min()
-            bundle.install_plugins()
-            logging.info(f"Installed plugins: {bundle.installed_plugins}")
+    with Bundles.create(build_manifest, artifacts_dir, bundle_recorder, args.keep) as bundle:
+        bundle.install_min()
+        bundle.install_plugins()
+        logging.info(f"Installed plugins: {bundle.installed_plugins}")
 
-            #  Save a copy of the manifest inside of the tar
-            bundle_recorder.write_manifest(bundle.min_dist.archive_path)
-            bundle.package(output_dir)
+        #  Save a copy of the manifest inside of the tar
+        bundle_recorder.write_manifest(bundle.min_dist.archive_path)
+        bundle.package(output_dir)
 
-            bundle_recorder.write_manifest(output_dir)
+        bundle_recorder.write_manifest(output_dir)
 
     logging.info("Done.")
 

tests/test_run_assemble.py

@@ -5,7 +5,6 @@
 # compatible open source license.
 
 import os
-import tempfile
 import unittest
 from unittest.mock import MagicMock, call, patch
 
@@ -31,14 +30,12 @@ def test_usage(self, *mocks):
 
     @patch("os.chdir")
     @patch("os.makedirs")
+    @patch("shutil.copy2")
     @patch("os.getcwd", return_value="curdir")
     @patch("argparse._sys.argv", ["run_assemble.py", BUILD_MANIFEST])
     @patch("run_assemble.Bundles.create")
     @patch("run_assemble.BundleRecorder", return_value=MagicMock())
-    @patch("run_assemble.TemporaryDirectory")
-    @patch("shutil.copy2")
-    def test_main(self, mock_copy, mock_temp, mock_recorder, mock_bundles, *mocks):
-        mock_temp.return_value.__enter__.return_value.name = tempfile.gettempdir()
+    def test_main(self, mock_recorder, mock_bundles, *mocks):
         mock_bundle = MagicMock(min_dist=MagicMock(archive_path="path"))
         mock_bundles.return_value.__enter__.return_value = mock_bundle
 

tests/tests_assemble_workflow/test_assemble_args.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+
+import logging
+import os
+import unittest
+from unittest.mock import patch
+
+from assemble_workflow.assemble_args import AssembleArgs
+
+
+class TestAssembleArgs(unittest.TestCase):
+
+    ASSEMBLE_PY = "./src/run_assembly.py"
+
+    OPENSEARCH_MANIFEST = os.path.realpath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "..",
+            "..",
+            "manifests",
+            "1.1.0",
+            "opensearch-1.1.0.yml",
+        )
+    )
+
+    @patch("argparse._sys.argv", [ASSEMBLE_PY, OPENSEARCH_MANIFEST])
+    def test_manifest(self):
+        self.assertEqual(AssembleArgs().manifest.name, TestAssembleArgs.OPENSEARCH_MANIFEST)
+
+    @patch("argparse._sys.argv", [ASSEMBLE_PY, OPENSEARCH_MANIFEST])
+    def test_keep_default(self):
+        self.assertFalse(AssembleArgs().keep)
+
+    @patch("argparse._sys.argv", [ASSEMBLE_PY, OPENSEARCH_MANIFEST, "--keep"])
+    def test_keep_true(self):
+        self.assertTrue(AssembleArgs().keep)
+
+    @patch("argparse._sys.argv", [ASSEMBLE_PY, OPENSEARCH_MANIFEST])
+    def test_verbose_default(self):
+        self.assertEqual(AssembleArgs().logging_level, logging.INFO)
+
+    @patch("argparse._sys.argv", [ASSEMBLE_PY, OPENSEARCH_MANIFEST, "--verbose"])
+    def test_verbose_true(self):
+        self.assertTrue(AssembleArgs().logging_level, logging.DEBUG)
+
+    @patch("argparse._sys.argv", [ASSEMBLE_PY, OPENSEARCH_MANIFEST, "--base-url", "url"])
+    def test_base_url(self):
+        self.assertEqual(AssembleArgs().base_url, "url")

tests/tests_assemble_workflow/test_bundles.py

@@ -18,14 +18,21 @@ class TestBundles(unittest.TestCase):
     def test_bundle_opensearch(self):
         manifest_path = os.path.join(os.path.dirname(__file__), "data", "opensearch-build-linux-1.1.0.yml")
         artifacts_path = os.path.join(os.path.dirname(__file__), "data", "artifacts")
-        bundle = Bundles.create(BuildManifest.from_path(manifest_path), artifacts_path, MagicMock())
+        bundle = Bundles.create(BuildManifest.from_path(manifest_path), artifacts_path, MagicMock(), False)
         self.assertIs(type(bundle), BundleOpenSearch)
 
     def test_bundle_opensearch_dashboards(self):
         manifest_path = os.path.join(os.path.dirname(__file__), "data", "opensearch-dashboards-build-1.1.0.yml")
         artifacts_path = os.path.join(os.path.dirname(__file__), "data", "artifacts")
-        bundle = Bundles.create(BuildManifest.from_path(manifest_path), artifacts_path, MagicMock())
+        bundle = Bundles.create(BuildManifest.from_path(manifest_path), artifacts_path, MagicMock(), False)
         self.assertIs(type(bundle), BundleOpenSearchDashboards)
+        self.assertFalse(bundle.tmp_dir.keep)
+
+    def test_bundle_keep(self):
+        manifest_path = os.path.join(os.path.dirname(__file__), "data", "opensearch-build-linux-1.1.0.yml")
+        artifacts_path = os.path.join(os.path.dirname(__file__), "data", "artifacts")
+        bundle = Bundles.create(BuildManifest.from_path(manifest_path), artifacts_path, MagicMock(), True)
+        self.assertTrue(bundle.tmp_dir.keep)
 
     def test_bundle_opensearch_invalid(self):
         manifest = BuildManifest(
@@ -41,5 +48,5 @@ def test_bundle_opensearch_invalid(self):
             }
         )
         with self.assertRaises(ValueError) as ctx:
-            Bundles.create(manifest, "path", MagicMock())
+            Bundles.create(manifest, "path", MagicMock(), False)
         self.assertEqual(str(ctx.exception), "Unsupported bundle: invalid")