load yaml tests from same branch as client

Author: miguelgrinberg

.buildkite/run-repository.sh

@@ -21,6 +21,7 @@ echo -e "\033[34;1mINFO:\033[0m TEST_SUITE ${TEST_SUITE}\033[0m"
 echo -e "\033[34;1mINFO:\033[0m NOX_SESSION ${NOX_SESSION}\033[0m"
 echo -e "\033[34;1mINFO:\033[0m PYTHON_VERSION ${PYTHON_VERSION}\033[0m"
 echo -e "\033[34;1mINFO:\033[0m PYTHON_CONNECTION_CLASS ${PYTHON_CONNECTION_CLASS}\033[0m"
+echo -e "\033[34;1mINFO:\033[0m YAML_TESTS_BRANCH ${YAML_TESTS_BRANCH}\033[0m"
 
 echo -e "\033[1m>>>>> Build [elastic/elasticsearch-py container] >>>>>>>>>>>>>>>>>>>>>>>>>>>>>\033[0m"
 

.buildkite/run-tests

@@ -9,6 +9,7 @@ export STACK_VERSION="${STACK_VERSION:=8.0.0-SNAPSHOT}"
 export TEST_SUITE="${TEST_SUITE:=platinum}"
 export PYTHON_VERSION="${PYTHON_VERSION:=3.14}"
 export PYTHON_CONNECTION_CLASS="${PYTHON_CONNECTION_CLASS:=urllib3}"
+export YAML_TESTS_BRANCH=${BUILDKITE_PULL_REQUEST_BASE_BRANCH:=main}
 
 script_path=$(dirname $(realpath $0))
 source $script_path/functions/imports.sh

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
@@ -35,6 +36,7 @@
 
 from elasticsearch import ApiError, ElasticsearchWarning, RequestError
 from elasticsearch._sync.client.utils import _base64_auth_header
+from elasticsearch._version import __versionstr__
 from elasticsearch.compat import string_types
 
 # some params had to be changed in python, keep track of them so we can rename
@@ -499,12 +501,28 @@ 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_candidates = []
+    if "YAML_TESTS_BRANCH" in os.environ:
+        branch_candidates.append(os.environ["YAML_TESTS_BRANCH"] or "main")
+    git_branch = subprocess.getoutput("git branch --show-current")
+    if git_branch not in branch_candidates:
+        branch_candidates.append(git_branch)
+    package_version = __versionstr__.rsplit(".", 1)[0]
+    if package_version not in branch_candidates:
+        branch_candidates.append(package_version)
+    if "main" not in branch_candidates:
+        branch_candidates.append("main")
+
+    response = None
+    for branch in branch_candidates:
+        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:
+            break
+    assert response is not None
 
     # 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 +567,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,