load yaml tests from same branch as client

Author: miguelgrinberg

test_elasticsearch/test_server/test_rest_api_spec.py

@@ -25,6 +25,7 @@
 import json
 import os
 import re
+import subprocess
 import warnings
 import zipfile
 from typing import Tuple, Union
@@ -499,12 +500,17 @@ def remove_implicit_resolver(cls, tag_to_remove):
     # Construct the HTTP and Elasticsearch client
     http = urllib3.PoolManager(retries=urllib3.Retry(total=10))
 
-    yaml_tests_url = (
-        "https://api.github.com/repos/elastic/elasticsearch-clients-tests/zipball/main"
-    )
+    branch = subprocess.getoutput("git branch --show-current")
+    yaml_tests_url = f"https://api.github.com/repos/elastic/elasticsearch-clients-tests/zipball/{branch}"
+    response = http.request("GET", yaml_tests_url)
+    if response.status == 404 and branch != "main":
+        # this branch does not exist in the yaml tests repo, so we use main instead
+        branch = "main"
+        yaml_tests_url = f"https://api.github.com/repos/elastic/elasticsearch-clients-tests/zipball/{branch}"
+        response = http.request("GET", yaml_tests_url)
 
     # Download the zip and start reading YAML from the files in memory
-    package_zip = zipfile.ZipFile(io.BytesIO(http.request("GET", yaml_tests_url).data))
+    package_zip = zipfile.ZipFile(io.BytesIO(response.data))
 
     for yaml_file in package_zip.namelist():
         if not re.match(r"^.*\/tests\/.*\.ya?ml$", yaml_file):
@@ -549,7 +555,7 @@ def remove_implicit_resolver(cls, tag_to_remove):
             for prefix in ("rest-api-spec/", "test/", "free/", "platinum/"):
                 if pytest_test_name.startswith(prefix):
                     pytest_test_name = pytest_test_name[len(prefix) :]
-            pytest_param_id = "%s[%d]" % (pytest_test_name, test_number)
+            pytest_param_id = "[%s]%s[%d]" % (branch, pytest_test_name, test_number)
 
             pytest_param = {
                 "setup": setup_steps,