From f973e405eeb4f28a6a937c26d19cad54acd00eb4 Mon Sep 17 00:00:00 2001
From: Max Kainov <max.kainov@clickhouse.com>
Date: Wed, 7 Feb 2024 19:36:28 +0000
Subject: [PATCH] CI: fixing ARM integration tests

 #do_not_test
---
 tests/ci/ci.py | 41 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/tests/ci/ci.py b/tests/ci/ci.py
index 9d57f161be3..898d23be843 100644
--- a/tests/ci/ci.py
+++ b/tests/ci/ci.py
@@ -1183,13 +1183,13 @@ def _configure_jobs(
 
         if batches_to_do:
             jobs_to_do.append(job)
+            jobs_params[job] = {
+                "batches": batches_to_do,
+                "num_batches": num_batches,
+            }
         elif add_to_skip:
             # treat job as being skipped only if it's controlled by digest
             jobs_to_skip.append(job)
-        jobs_params[job] = {
-            "batches": batches_to_do,
-            "num_batches": num_batches,
-        }
 
     if not pr_info.is_release_branch():
         # randomization bucket filtering (pick one random job from each bucket, for jobs with configured random_bucket property)
@@ -1268,6 +1268,33 @@ def _configure_jobs(
             jobs_to_do = list(
                 set(job for job in jobs_to_do_requested if job not in jobs_to_skip)
             )
+            # if requested job does not have params in jobs_params (it happens for "run_by_label" job)
+            #   we need to add params - otherwise it won't run as "batches" list will be empty
+            for job in jobs_to_do:
+                if job not in jobs_params:
+                    num_batches = CI_CONFIG.get_job_config(job).num_batches
+                    jobs_params[job] = {
+                        "batches": list(range(num_batches)),
+                        "num_batches": num_batches,
+                    }
+
+        requested_batches = set()
+        for token in commit_tokens:
+            if token.startswith("batch_"):
+                try:
+                    batches = [
+                        int(batch) for batch in token.removeprefix("batch_").split("_")
+                    ]
+                except Exception:
+                    print(f"ERROR: failed to parse commit tag [{token}]")
+                requested_batches.update(batches)
+        if requested_batches:
+            print(
+                f"NOTE: Only specific job batches were requested [{list(requested_batches)}]"
+            )
+            for job, params in jobs_params.items():
+                if params["num_batches"] > 1:
+                    params["batches"] = list(requested_batches)
 
     return {
         "digests": digests,
@@ -1372,7 +1399,11 @@ def _update_gh_statuses_action(indata: Dict, s3: S3Helper) -> None:
 def _fetch_commit_tokens(message: str) -> List[str]:
     pattern = r"#[\w-]+"
     matches = [match[1:] for match in re.findall(pattern, message)]
-    res = [match for match in matches if match in Labels or match.startswith("job_")]
+    res = [
+        match
+        for match in matches
+        if match in Labels or match.startswith("job_") or match.startswith("batch_")
+    ]
     return res