From fc211dee6614f4e546e08b534b9bf4004aab2b8d Mon Sep 17 00:00:00 2001
From: Dennis Lawler <4824647+drawlerr@users.noreply.github.com>
Date: Wed, 2 Oct 2019 06:45:06 -0700
Subject: [PATCH] Add support for OSNAME and ARCH variables in dist repo URLs.
 (#781)

---
 esrally/mechanic/supplier.py    | 16 ++++++++++++++--
 esrally/utils/sysstats.py       |  7 +++++++
 tests/mechanic/supplier_test.py | 12 +++++++-----
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/esrally/mechanic/supplier.py b/esrally/mechanic/supplier.py
index 58173cc8b..daa816a42 100644
--- a/esrally/mechanic/supplier.py
+++ b/esrally/mechanic/supplier.py
@@ -23,7 +23,7 @@
 
 from esrally import exceptions, PROGRAM_NAME
 from esrally.exceptions import BuildError, SystemSetupError
-from esrally.utils import git, io, process, net, jvm, convert
+from esrally.utils import git, io, process, net, jvm, convert, sysstats
 
 # e.g. my-plugin:current - we cannot simply use String#split(":") as this would not work for timestamp-based revisions
 REVISION_PATTERN = r"(\w.*?):(.*)"
@@ -525,7 +525,19 @@ def _url_for(self, user_defined_key, default_key, mandatory=True):
                 raise exceptions.SystemSetupError("Neither config key [{}] nor [{}] is defined.".format(user_defined_key, default_key))
             else:
                 return None
-        return url_template.replace("{{VERSION}}", self.version)
+        return self._substitute_vars(url_template)
+
+    def _substitute_vars(self, s):
+        substitutions = {
+            "{{VERSION}}": self.version,
+            "{{OSNAME}}": sysstats.os_name().lower(),
+            "{{ARCH}}": sysstats.cpu_arch().lower()
+        }
+        r = s
+        for key, replacement in substitutions.items():
+            r = r.replace(key, replacement)
+        return r
+
 
     @property
     def cache(self):
diff --git a/esrally/utils/sysstats.py b/esrally/utils/sysstats.py
index 66a9d44cf..456814d21 100644
--- a/esrally/utils/sysstats.py
+++ b/esrally/utils/sysstats.py
@@ -53,6 +53,13 @@ def cpu_model():
     return "Unknown"
 
 
+def cpu_arch():
+    """
+    :return: The CPU architecture name.
+    """
+    return platform.uname().machine
+
+
 def disks():
     # only physical partitions (no memory partitions etc)
     return psutil.disk_partitions(all=False)
diff --git a/tests/mechanic/supplier_test.py b/tests/mechanic/supplier_test.py
index 199a91368..7b61bb3fe 100644
--- a/tests/mechanic/supplier_test.py
+++ b/tests/mechanic/supplier_test.py
@@ -523,13 +523,15 @@ def test_create_suppliers_for_es_and_plugin_source_build(self):
 
 
 class DistributionRepositoryTests(TestCase):
-    def test_release_repo_config_with_default_url(self):
+    @mock.patch("esrally.utils.sysstats.os_name", return_value="Linux")
+    @mock.patch("esrally.utils.sysstats.cpu_arch", return_value="X86_64")
+    def test_release_repo_config_with_default_url(self, os_name, cpu_arch):
         repo = supplier.DistributionRepository(name="release", distribution_config={
-            "release_url": "https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-{{VERSION}}.tar.gz",
+            "release_url": "https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-{{VERSION}}-{{OSNAME}}-{{ARCH}}.tar.gz",
             "release.cache": "true"
-        }, version="5.5.0")
-        self.assertEqual("https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-5.5.0.tar.gz", repo.download_url)
-        self.assertEqual("elasticsearch-5.5.0.tar.gz", repo.file_name)
+        }, version="7.3.2")
+        self.assertEqual("https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.3.2-linux-x86_64.tar.gz", repo.download_url)
+        self.assertEqual("elasticsearch-7.3.2-linux-x86_64.tar.gz", repo.file_name)
         self.assertTrue(repo.cache)
 
     def test_release_repo_config_with_user_url(self):