Refactors for using BinaryKey enum, factory pattern, and creating Art…

…ifactVariablesProvider helper class (#183) * Refactors for using BinaryKey enum, factory pattern, and creating ArtifactVariablesProvider helper class Signed-off-by: Chase Engelbrecht <[email protected]> * Add unit tests for ArtifactVariablesProvider Signed-off-by: Chase Engelbrecht <[email protected]> * Refactor ArchitectureTypes to an enum Signed-off-by: Chase Engelbrecht <[email protected]>
opensearch-project · Apr 13, 2022 · d11db37 · d11db37
1 parent 4dbaf99
commit d11db37
Show file tree

Hide file tree

Showing 16 changed files with 142 additions and 77 deletions.
diff --git a/osbenchmark/builder/downloaders/opensearch_distribution_downloader.py b/osbenchmark/builder/downloaders/opensearch_distribution_downloader.py
@@ -2,25 +2,21 @@
 import os.path
 
 from osbenchmark.builder.downloaders.downloader import Downloader
-from osbenchmark.builder.downloaders.repositories.opensearch_distribution_repository_provider import \
-    OpenSearchDistributionRepositoryProvider
-from osbenchmark.builder.utils.path_manager import PathManager
+from osbenchmark.builder.utils.binary_keys import BinaryKeys
 from osbenchmark.exceptions import ExecutorError
 
 
 class OpenSearchDistributionDownloader(Downloader):
-    BINARY_KEY = "opensearch"
-
-    def __init__(self, provision_config_instance, executor):
+    def __init__(self, provision_config_instance, executor, path_manager, distribution_repository_provider):
         super().__init__(executor)
         self.logger = logging.getLogger(__name__)
         self.provision_config_instance = provision_config_instance
-        self.path_manager = PathManager(executor)
-        self.distribution_repository_provider = OpenSearchDistributionRepositoryProvider(provision_config_instance, executor)
+        self.path_manager = path_manager
+        self.distribution_repository_provider = distribution_repository_provider
 
     def download(self, host):
         binary_path = self._fetch_binary(host)
-        return {OpenSearchDistributionDownloader.BINARY_KEY: binary_path}
+        return {BinaryKeys.OPENSEARCH: binary_path}
 
     def _fetch_binary(self, host):
         download_url = self.distribution_repository_provider.get_download_url(host)

diff --git a/osbenchmark/builder/downloaders/plugin_distribution_downloader.py b/osbenchmark/builder/downloaders/plugin_distribution_downloader.py
@@ -1,13 +1,11 @@
 from osbenchmark.builder.downloaders.downloader import Downloader
-from osbenchmark.builder.downloaders.repositories.plugin_distribution_repository_provider import \
-    PluginDistributionRepositoryProvider
 
 
 class PluginDistributionDownloader(Downloader):
-    def __init__(self, plugin, executor):
+    def __init__(self, plugin, executor, distribution_repository_provider):
         super().__init__(executor)
         self.plugin = plugin
-        self.distribution_repository_provider = PluginDistributionRepositoryProvider(plugin, executor)
+        self.distribution_repository_provider = distribution_repository_provider
 
     def download(self, host):
         plugin_url = self.distribution_repository_provider.get_download_url(host)

diff --git a/osbenchmark/builder/downloaders/repositories/opensearch_distribution_repository_provider.py b/osbenchmark/builder/downloaders/repositories/opensearch_distribution_repository_provider.py
@@ -1,15 +1,13 @@
 import logging
 
-from osbenchmark.builder.downloaders.repositories.repository_url_provider import RepositoryUrlProvider
 from osbenchmark.utils import convert
 
 
 class OpenSearchDistributionRepositoryProvider:
-    def __init__(self, provision_config_instance, executor):
+    def __init__(self, provision_config_instance, repository_url_provider):
         self.logger = logging.getLogger(__name__)
         self.provision_config_instance = provision_config_instance
-        self.executor = executor
-        self.repository_url_provider = RepositoryUrlProvider(executor)
+        self.repository_url_provider = repository_url_provider
 
     def get_download_url(self, host):
         is_runtime_jdk_bundled = self.provision_config_instance.variables["system"]["runtime"]["jdk"]["bundled"]

diff --git a/osbenchmark/builder/downloaders/repositories/plugin_distribution_repository_provider.py b/osbenchmark/builder/downloaders/repositories/plugin_distribution_repository_provider.py
@@ -1,10 +1,7 @@
-from osbenchmark.builder.downloaders.repositories.repository_url_provider import RepositoryUrlProvider
-
-
 class PluginDistributionRepositoryProvider:
-    def __init__(self, plugin, executor):
+    def __init__(self, plugin, repository_url_provider):
         self.plugin = plugin
-        self.repository_url_provider = RepositoryUrlProvider(executor)
+        self.repository_url_provider = repository_url_provider
 
     def get_download_url(self, host):
         distribution_repository = self.plugin.variables["distribution"]["repository"]

diff --git a/osbenchmark/builder/downloaders/repositories/repository_url_provider.py b/osbenchmark/builder/downloaders/repositories/repository_url_provider.py
@@ -1,18 +1,12 @@
 from functools import reduce
 
-from osbenchmark.builder.utils.template_renderer import TemplateRenderer
 from osbenchmark.exceptions import SystemSetupError
 
-ARCH_MAPPINGS = {
-    "x86_64": "x64",
-    "aarch64": "arm64"
-}
-
 
 class RepositoryUrlProvider:
-    def __init__(self, executor):
-        self.executor = executor
-        self.template_renderer = TemplateRenderer()
+    def __init__(self, template_renderer, artifact_variables_provider):
+        self.template_renderer = template_renderer
+        self.artifact_variables_provider = artifact_variables_provider
 
     def render_url_for_key(self, host, config_variables, key, mandatory=True):
         try:
@@ -22,22 +16,10 @@ def render_url_for_key(self, host, config_variables, key, mandatory=True):
                 raise SystemSetupError(f"Config key [{key}] is not defined.")
             else:
                 return None
-        return self.template_renderer.render_template_string(url_template, self._get_url_template_variables(host, config_variables))
+
+        artifact_version = config_variables["distribution"]["version"]
+        artifact_variables = self.artifact_variables_provider.get_artifact_variables(host, artifact_version)
+        return self.template_renderer.render_template_string(url_template, artifact_variables)
 
     def _get_value_from_dot_notation_key(self, dict_object, key):
         return reduce(dict.get, key.split("."), dict_object)
-
-    def _get_url_template_variables(self, host, config_variables):
-        return {
-            "VERSION": config_variables["distribution"]["version"],
-            "OSNAME": self._get_os_name(host),
-            "ARCH": self._get_arch(host)
-        }
-
-    def _get_os_name(self, host):
-        os_name = self.executor.execute(host, "uname", output=True)[0]
-        return os_name.lower()
-
-    def _get_arch(self, host):
-        arch = self.executor.execute(host, "uname -m", output=True)[0]
-        return ARCH_MAPPINGS[arch.lower()]
diff --git a/osbenchmark/builder/installers/preparers/opensearch_preparer.py b/osbenchmark/builder/installers/preparers/opensearch_preparer.py
@@ -4,13 +4,12 @@
 
 from osbenchmark.builder.installers.preparers.preparer import Preparer
 from osbenchmark.builder.models.node import Node
+from osbenchmark.builder.utils.binary_keys import BinaryKeys
 from osbenchmark.builder.utils.host_cleaner import HostCleaner
 from osbenchmark.builder.utils.path_manager import PathManager
 
 
 class OpenSearchPreparer(Preparer):
-    OPENSEARCH_BINARY_KEY = "opensearch"
-
     def __init__(self, provision_config_instance, executor, hook_handler_class):
         super().__init__(executor)
         self.logger = logging.getLogger(__name__)
@@ -23,7 +22,7 @@ def __init__(self, provision_config_instance, executor, hook_handler_class):
 
     def prepare(self, host, binaries):
         node = self._create_node()
-        self._prepare_node(host, node, binaries[OpenSearchPreparer.OPENSEARCH_BINARY_KEY])
+        self._prepare_node(host, node, binaries[BinaryKeys.OPENSEARCH])
 
         return node
 

diff --git a/osbenchmark/builder/models/architecture_types.py b/osbenchmark/builder/models/architecture_types.py
@@ -0,0 +1,25 @@
+from enum import Enum
+
+
+class ArchitectureTypes(Enum):
+    """
+    Represents a machine's architecture type
+
+    :param hardware_name: The value returned by the machine when querying the architecture. Obtained via `uname -m` for unix machines
+    ;param opensearch_name: The value used by opensearch artifacts to represent the architecture
+    """
+
+    def __init__(self, hardware_name, opensearch_name):
+        self.hardware_name = hardware_name
+        self.opensearch_name = opensearch_name
+
+    ARM = "aarch64", "arm64"
+    x86 = "x86_64", "x64"
+
+    @staticmethod
+    def get_from_hardware_name(hardware_name):
+        for arch_type in ArchitectureTypes:
+            if arch_type.hardware_name == hardware_name:
+                return arch_type
+
+        raise ValueError
diff --git a/osbenchmark/builder/utils/artifact_variables_provider.py b/osbenchmark/builder/utils/artifact_variables_provider.py
@@ -0,0 +1,21 @@
+from osbenchmark.builder.models.architecture_types import ArchitectureTypes
+
+
+class ArtifactVariablesProvider:
+    def __init__(self, executor):
+        self.executor = executor
+
+    def get_artifact_variables(self, host, opensearch_version=None):
+        return {
+            "VERSION": opensearch_version,
+            "OSNAME": self._get_os_name(host),
+            "ARCH": self._get_arch(host)
+        }
+
+    def _get_os_name(self, host):
+        os_name = self.executor.execute(host, "uname", output=True)[0]
+        return os_name.lower()
+
+    def _get_arch(self, host):
+        arch = self.executor.execute(host, "uname -m", output=True)[0]
+        return ArchitectureTypes.get_from_hardware_name(arch.lower()).opensearch_name
diff --git a/osbenchmark/builder/utils/binary_keys.py b/osbenchmark/builder/utils/binary_keys.py
@@ -0,0 +1,5 @@
+from enum import Enum
+
+
+class BinaryKeys(str, Enum):
+    OPENSEARCH = "opensearch"
diff --git a/tests/builder/downloaders/opensearch_distribution_downloader_test.py b/tests/builder/downloaders/opensearch_distribution_downloader_test.py
@@ -22,9 +22,11 @@ def setUp(self):
             }
         })
 
-        self.os_distro_downloader = OpenSearchDistributionDownloader(self.provision_config_instance, self.executor)
-        self.os_distro_downloader.path_manager = Mock()
-        self.os_distro_downloader.distribution_repository_provider = Mock()
+        self.path_manager = Mock()
+        self.distribution_repository_provider = Mock()
+        self.os_distro_downloader = OpenSearchDistributionDownloader(self.provision_config_instance, self.executor, self.path_manager,
+                                                                     self.distribution_repository_provider)
+
 
         self.os_distro_downloader.distribution_repository_provider.get_download_url.return_value = "https://fake/download.tar.gz"
         self.os_distro_downloader.distribution_repository_provider.get_file_name_from_download_url.return_value = "my-distro"

diff --git a/tests/builder/downloaders/plugin_distribution_downloader_test.py b/tests/builder/downloaders/plugin_distribution_downloader_test.py
@@ -12,8 +12,8 @@ def setUp(self):
         self.executor = Mock()
         self.plugin = PluginDescriptor(name="my plugin")
 
-        self.plugin_distro_downloader = PluginDistributionDownloader(self.plugin, self.executor)
-        self.plugin_distro_downloader.distribution_repository_provider = Mock()
+        self.distribution_repository_provider = Mock()
+        self.plugin_distro_downloader = PluginDistributionDownloader(self.plugin, self.executor, self.distribution_repository_provider)
 
     def test_plugin_url_exists(self):
         self.plugin_distro_downloader.distribution_repository_provider.get_download_url.return_value = "https://fake"

diff --git a/tests/builder/downloaders/repositories/opensearch_distribution_repository_provider_test.py b/tests/builder/downloaders/repositories/opensearch_distribution_repository_provider_test.py
@@ -8,8 +8,6 @@
 
 class OpenSearchDistributionRepositoryProviderTest(TestCase):
     def setUp(self):
-        self.executor = Mock()
-
         self.host = None
         self.provision_config_instance = ProvisionConfigInstance(names=None, config_paths=None, root_path=None, variables={
             "system": {
@@ -26,8 +24,10 @@ def setUp(self):
                 }
             }
         })
-        self.os_distro_repo_provider = OpenSearchDistributionRepositoryProvider(self.provision_config_instance, self.executor)
-        self.os_distro_repo_provider.repository_url_provider = Mock()
+        self.repository_url_provider = Mock()
+        self.os_distro_repo_provider = OpenSearchDistributionRepositoryProvider(self.provision_config_instance,
+                                                                                self.repository_url_provider)
+
 
     def test_get_url_bundled_jdk(self):
         self.os_distro_repo_provider.get_download_url(self.host)

diff --git a/tests/builder/downloaders/repositories/plugin_distribution_repository_provider_test.py b/tests/builder/downloaders/repositories/plugin_distribution_repository_provider_test.py
@@ -8,12 +8,11 @@
 
 class PluginDistributionRepositoryProviderTest(TestCase):
     def setUp(self):
-        self.executor = Mock()
-
         self.host = None
         self.plugin = PluginDescriptor(name="my-plugin", variables={"distribution": {"repository": "release"}})
-        self.plugin_distro_repo_provider = PluginDistributionRepositoryProvider(self.plugin, self.executor)
-        self.plugin_distro_repo_provider.repository_url_provider = Mock()
+        self.repository_url_provider = Mock()
+        self.plugin_distro_repo_provider = PluginDistributionRepositoryProvider(self.plugin, self.repository_url_provider)
+
 
     def test_get_plugin_url(self):
         self.plugin_distro_repo_provider.get_download_url(self.host)

diff --git a/tests/builder/downloaders/repositories/repository_url_provider_test.py b/tests/builder/downloaders/repositories/repository_url_provider_test.py
@@ -1,4 +1,4 @@
-from unittest import TestCase
+from unittest import TestCase, mock
 from unittest.mock import Mock
 
 from osbenchmark.builder.downloaders.repositories.repository_url_provider import RepositoryUrlProvider
@@ -7,7 +7,8 @@
 
 class RepositoryUrlProviderTest(TestCase):
     def setUp(self):
-        self.executor = Mock()
+        self.template_renderer = Mock()
+        self.artifact_variables_provider = Mock()
 
         self.host = None
         self.variables = {
@@ -20,19 +21,18 @@ def setUp(self):
         }
         self.url_key = "fake.url"
 
-        self.repo_url_provider = RepositoryUrlProvider(self.executor)
+        self.repo_url_provider = RepositoryUrlProvider(self.template_renderer, self.artifact_variables_provider)
 
-    def test_get_url_aarch64(self):
-        self.executor.execute.side_effect = [["Linux"], ["aarch64"]]
+    def test_get_url(self):
+        self.artifact_variables_provider.get_artifact_variables.return_value = {"fake": "vars"}
 
-        url = self.repo_url_provider.render_url_for_key(self.host, self.variables, self.url_key)
-        self.assertEqual(url, "opensearch/1.2.3/opensearch-1.2.3-linux-arm64.tar.gz")
-
-    def test_get_url_x86(self):
-        self.executor.execute.side_effect = [["Linux"], ["x86_64"]]
-
-        url = self.repo_url_provider.render_url_for_key(self.host, self.variables, self.url_key)
-        self.assertEqual(url, "opensearch/1.2.3/opensearch-1.2.3-linux-x64.tar.gz")
+        self.repo_url_provider.render_url_for_key(self.host, self.variables, self.url_key)
+        self.artifact_variables_provider.get_artifact_variables.assert_has_calls([
+            mock.call(self.host, "1.2.3")
+        ])
+        self.template_renderer.render_template_string.assert_has_calls([
+            mock.call("opensearch/{{VERSION}}/opensearch-{{VERSION}}-{{OSNAME}}-{{ARCH}}.tar.gz", {"fake": "vars"})
+        ])
 
     def test_no_url_template_found(self):
         with self.assertRaises(SystemSetupError):

diff --git a/tests/builder/installers/preparers/opensearch_preparer_test.py b/tests/builder/installers/preparers/opensearch_preparer_test.py
@@ -6,6 +6,7 @@
 from osbenchmark.builder.models.host import Host
 from osbenchmark.builder.models.node import Node
 from osbenchmark.builder.provision_config import ProvisionConfigInstance
+from osbenchmark.builder.utils.binary_keys import BinaryKeys
 
 
 class OpenSearchPreparerTests(TestCase):
@@ -15,7 +16,7 @@ def setUp(self):
                          name=self.node_id, pid=None, telemetry=None, port=9200, root_dir=None,
                          log_path="/fake/logpath", heap_dump_path="/fake/heap")
         self.host = Host(name="fake", address="10.17.22.23", metadata={}, node=None)
-        self.binaries = {OpenSearchPreparer.OPENSEARCH_BINARY_KEY: "/data/builds/distributions"}
+        self.binaries = {BinaryKeys.OPENSEARCH: "/data/builds/distributions"}
         self.all_node_ips = ["10.17.22.22", "10.17.22.23"]
 
         self.test_execution_root = "fake_root"

diff --git a/tests/builder/utils/artifact_variables_provider_test.py b/tests/builder/utils/artifact_variables_provider_test.py
@@ -0,0 +1,42 @@
+from unittest import TestCase
+from unittest.mock import Mock
+
+from osbenchmark.builder.utils.artifact_variables_provider import ArtifactVariablesProvider
+
+
+class ArtifactVariablesProviderTest(TestCase):
+    def setUp(self):
+        self.host = None
+
+        self.executor = Mock()
+        self.artifact_variables_provider = ArtifactVariablesProvider(self.executor)
+
+    def test_x86(self):
+        self.executor.execute.side_effect = [["Linux"], ["x86_64"]]
+        variables = self.artifact_variables_provider.get_artifact_variables(self.host)
+
+        self.assertEqual(variables, {
+            "VERSION": None,
+            "OSNAME": "linux",
+            "ARCH": "x64"
+        })
+
+    def test_arm(self):
+        self.executor.execute.side_effect = [["Linux"], ["aarch64"]]
+        variables = self.artifact_variables_provider.get_artifact_variables(self.host)
+
+        self.assertEqual(variables, {
+            "VERSION": None,
+            "OSNAME": "linux",
+            "ARCH": "arm64"
+        })
+
+    def test_version_supplied(self):
+        self.executor.execute.side_effect = [["Linux"], ["aarch64"]]
+        variables = self.artifact_variables_provider.get_artifact_variables(self.host, "1.23")
+
+        self.assertEqual(variables, {
+            "VERSION": "1.23",
+            "OSNAME": "linux",
+            "ARCH": "arm64"
+        })