From 4dd5e4fe993b8ca2a5f693f8f809fa0c2fc76fa2 Mon Sep 17 00:00:00 2001
From: Max Kainov <max.kainov@clickhouse.com>
Date: Wed, 6 Mar 2024 16:30:47 +0000
Subject: [PATCH] CI: fix stage config (unit test release issue)

 #no_ci_cache
---
 tests/ci/ci_config.py      | 16 ++++++++++++----
 tests/ci/test_ci_config.py | 20 +++++++++++++++++++-
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/tests/ci/ci_config.py b/tests/ci/ci_config.py
index 308a9098c29..4d944d24765 100644
--- a/tests/ci/ci_config.py
+++ b/tests/ci/ci_config.py
@@ -11,7 +11,7 @@ from ci_utils import WithIter
 from integration_test_images import IMAGES
 
 
-class CIStages:
+class CIStages(metaclass=WithIter):
     NA = "UNKNOWN"
     BUILDS_1 = "Builds_1"
     BUILDS_2 = "Builds_2"
@@ -547,9 +547,17 @@ class CIConfig:
             stage_type = CIStages.TESTS_2
         elif self.is_test_job(job_name):
             stage_type = CIStages.TESTS_1
-            if job_name == JobNames.LIBFUZZER_TEST:
-                # since fuzzers build in Builds_2, test must be in Tests_2
-                stage_type = CIStages.TESTS_2
+            if job_name in CI_CONFIG.test_configs:
+                required_build = CI_CONFIG.test_configs[job_name].required_build
+                assert required_build
+                if required_build in CI_CONFIG.get_builds_for_report(
+                    JobNames.BUILD_CHECK
+                ):
+                    stage_type = CIStages.TESTS_1
+                else:
+                    stage_type = CIStages.TESTS_2
+            else:
+                stage_type = CIStages.TESTS_1
         assert stage_type, f"BUG [{job_name}]"
         return stage_type
 
diff --git a/tests/ci/test_ci_config.py b/tests/ci/test_ci_config.py
index 04c90105276..badbc4c5dcf 100644
--- a/tests/ci/test_ci_config.py
+++ b/tests/ci/test_ci_config.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 import unittest
-from ci_config import JobNames, CI_CONFIG, Runners
+from ci_config import CIStages, JobNames, CI_CONFIG, Runners
 
 
 class TestCIConfig(unittest.TestCase):
@@ -10,3 +10,21 @@ class TestCIConfig(unittest.TestCase):
         for job in JobNames:
             runner = CI_CONFIG.get_runner_type(job)
             self.assertIn(runner, Runners)
+
+    def test_job_stage_config(self):
+        """check runner is provided w/o exception"""
+        for job in JobNames:
+            stage = CI_CONFIG.get_job_ci_stage(job)
+            if job in [
+                JobNames.STYLE_CHECK,
+                JobNames.FAST_TEST,
+                JobNames.JEPSEN_KEEPER,
+                JobNames.BUILD_CHECK,
+                JobNames.BUILD_CHECK_SPECIAL,
+            ]:
+                assert (
+                    stage == CIStages.NA
+                ), "These jobs are not in CI stages, must be NA"
+            else:
+                assert stage != CIStages.NA, f"stage not found for [{job}]"
+            self.assertIn(stage, CIStages)