[ci] Balance integration test buckets by recorded durations (#18895)

tests/script/test_determine_jobs.py

@@ -151,9 +151,14 @@ def test_main_all_tests_should_run(
         patch.object(determine_jobs, "_is_clang_tidy_full_scan", return_value=False),
         patch.object(
             determine_jobs,
-            "_all_integration_test_files",
+            "all_integration_test_files",
             return_value=fake_test_files,
         ),
+        patch.object(
+            determine_jobs,
+            "load_integration_durations",
+            return_value=dict.fromkeys(fake_test_files, 200.0),
+        ),
         patch.object(
             determine_jobs,
             "get_changed_components",
@@ -189,24 +194,12 @@ def test_main_all_tests_should_run(
     output = json.loads(captured.out)
 
     assert output["integration_tests"] is True
-    # run_all=True expands to the full glob and pre-buckets into 3 parts.
-    # Each bucket's `tests` is a JSON list of file paths.
+    assert output["integration_run_all"] is True
+    # run_all=True expands to the full glob; balance and naming are pinned
+    # by the unit tests, main() only needs to round-trip the structure
     assert isinstance(output["integration_test_buckets"], list)
-    assert len(output["integration_test_buckets"]) == 3
-    assert [b["name"] for b in output["integration_test_buckets"]] == [
-        "1/3",
-        "2/3",
-        "3/3",
-    ]
-    for bucket in output["integration_test_buckets"]:
-        assert isinstance(bucket["tests"], list)
-        for path in bucket["tests"]:
-            assert isinstance(path, str)
     bucket_files = [f for b in output["integration_test_buckets"] for f in b["tests"]]
-    assert bucket_files == fake_test_files
-    # Bucket sizes are balanced (max-min difference at most 1).
-    sizes = [len(b["tests"]) for b in output["integration_test_buckets"]]
-    assert max(sizes) - min(sizes) <= 1
+    assert sorted(bucket_files) == fake_test_files
     assert output["clang_tidy"] is True
     assert output["clang_tidy_mode"] in ["nosplit", "split"]
     assert output["clang_format"] is True
@@ -509,14 +502,24 @@ def test_compute_integration_test_buckets_at_threshold_stays_single() -> None:
 
 
 def test_compute_integration_test_buckets_just_over_threshold_splits() -> None:
-    """One file over the threshold triggers the 3-bucket fan-out, balanced."""
+    """One file over the threshold fans out fully when the weights demand it."""
     n = determine_jobs.INTEGRATION_TESTS_SPLIT_THRESHOLD + 1
     files = [f"tests/integration/test_{i:02d}.py" for i in range(n)]
-    run, buckets = determine_jobs._compute_integration_test_buckets(False, files)
+    with patch.object(
+        determine_jobs,
+        "load_integration_durations",
+        return_value=dict.fromkeys(files, 200.0),
+    ):
+        run, buckets = determine_jobs._compute_integration_test_buckets(False, files)
     assert run is True
-    assert [b["name"] for b in buckets] == ["1/3", "2/3", "3/3"]
-    union = [path for b in buckets for path in b["tests"]]
+    # threshold+1 files x 200s caps at the maximum bucket count.
+    n_buckets = determine_jobs.INTEGRATION_TESTS_SPLIT_BUCKETS
+    assert [b["name"] for b in buckets] == [
+        f"{i + 1}/{n_buckets}" for i in range(n_buckets)
+    ]
+    union = sorted(path for b in buckets for path in b["tests"])
     assert union == sorted(files)
+    # Equal weights => bucket sizes are balanced (difference at most 1).
     sizes = [len(b["tests"]) for b in buckets]
     assert max(sizes) - min(sizes) <= 1
 
@@ -526,7 +529,7 @@ def test_compute_integration_test_buckets_run_all_with_empty_glob_disables_run()
 ):
     """run_all=True but glob returns no files => run suppressed (otherwise
     pytest would collect tests outside tests/integration/)."""
-    with patch.object(determine_jobs, "_all_integration_test_files", return_value=[]):
+    with patch.object(determine_jobs, "all_integration_test_files", return_value=[]):
         run, buckets = determine_jobs._compute_integration_test_buckets(True, [])
     assert run is False
     assert buckets == []
@@ -3146,3 +3149,86 @@ def test_memory_impact_elf_layouts_are_found(tmp_path: Path) -> None:
         elf.write_text("")
 
         assert find_elf_path(build_path) == elf, f"{platform} ELF not found"
+
+
+def test_compute_integration_test_buckets_no_durations_full_fanout() -> None:
+    """Without recorded durations the fan-out stays at the maximum."""
+    files = [f"tests/integration/test_{i:03d}.py" for i in range(15)]
+    with patch.object(determine_jobs, "load_integration_durations", return_value={}):
+        run, buckets = determine_jobs._compute_integration_test_buckets(False, files)
+    assert run is True
+    assert len(buckets) == determine_jobs.INTEGRATION_TESTS_SPLIT_BUCKETS
+    assert sorted(f for b in buckets for f in b["tests"]) == files
+
+
+def test_compute_integration_test_buckets_adaptive_count() -> None:
+    """A small recorded total weight collapses to one bucket above the threshold."""
+    files = [f"tests/integration/test_{i:03d}.py" for i in range(15)]
+    with patch.object(
+        determine_jobs,
+        "load_integration_durations",
+        return_value=dict.fromkeys(files, 10.0),
+    ):
+        run, buckets = determine_jobs._compute_integration_test_buckets(False, files)
+    assert run is True
+    # 15 files x 10s recorded = 150s, under the per-bucket weight target.
+    assert [b["name"] for b in buckets] == ["1/1"]
+    assert buckets[0]["tests"] == files
+
+
+def test_compute_integration_test_buckets_duration_weighted() -> None:
+    """Heavy files spread across buckets instead of clustering by sorted name."""
+    files = [f"tests/integration/test_{i:03d}.py" for i in range(12)]
+    durations = dict.fromkeys(files, 10.0)
+    durations[files[0]] = 600.0
+    durations[files[1]] = 600.0
+    with patch.object(
+        determine_jobs, "load_integration_durations", return_value=durations
+    ):
+        run, buckets = determine_jobs._compute_integration_test_buckets(False, files)
+    assert run is True
+    assert len(buckets) >= 2
+    heavy_buckets = [b for b in buckets if set(files[:2]) & set(b["tests"])]
+    assert len(heavy_buckets) == 2, "heavy files should land in different buckets"
+    assert sorted(f for b in buckets for f in b["tests"]) == files
+
+
+def test_load_integration_durations_missing_or_corrupt(tmp_path: Path) -> None:
+    """Missing or unparsable durations data degrades to an empty mapping."""
+    with patch.object(helpers, "root_path", str(tmp_path)):
+        assert determine_jobs.load_integration_durations() == {}
+        durations_file = tmp_path / helpers.INTEGRATION_TEST_DURATIONS_FILE
+        durations_file.parent.mkdir(parents=True)
+        durations_file.write_text("not json")
+        assert determine_jobs.load_integration_durations() == {}
+        durations_file.write_text('{"tests/integration/test_a.py": 12.5}')
+        assert determine_jobs.load_integration_durations() == {
+            "tests/integration/test_a.py": 12.5
+        }
+        # Non-positive entries are dropped, valid ones survive
+        durations_file.write_text(
+            '{"tests/integration/test_a.py": 12.5, "tests/integration/test_b.py": -1}'
+        )
+        assert determine_jobs.load_integration_durations() == {
+            "tests/integration/test_a.py": 12.5
+        }
+        # One non-numeric entry cannot discard the whole recording
+        durations_file.write_text(
+            '{"tests/integration/test_a.py": 12.5, "tests/integration/test_b.py": null}'
+        )
+        assert determine_jobs.load_integration_durations() == {
+            "tests/integration/test_a.py": 12.5
+        }
+        # A non-dict top level degrades to empty
+        durations_file.write_text("[12.5]")
+        assert determine_jobs.load_integration_durations() == {}
+
+
+def test_committed_integration_durations_are_sane() -> None:
+    """The committed recording itself holds positive bounded floats."""
+    raw = json.loads(
+        (Path(helpers.root_path) / helpers.INTEGRATION_TEST_DURATIONS_FILE).read_text()
+    )
+    assert raw, "committed durations file missing or empty"
+    assert all(isinstance(v, (int, float)) and 0 < v < 86400 for v in raw.values())
+    assert all(k.startswith("tests/integration/test_") for k in raw)

tests/script/test_helpers.py

@@ -2120,3 +2120,30 @@ def test_get_cpp_changed_components_independent_of_cwd(
     assert helpers.get_cpp_changed_components(
         ["tests/components/time/__init__.py"]
     ) == ["time"]
+
+
+def test_lpt_partition_balances_skewed_weights() -> None:
+    """Heavy items spread across groups instead of clustering."""
+    items = [f"i{n}" for n in range(6)]
+    weights = {"i0": 100.0, "i1": 90.0, "i2": 10.0, "i3": 10.0, "i4": 5.0, "i5": 5.0}
+    groups = helpers.lpt_partition(items, weights, 2)
+    group_weights = sorted(sum(weights[i] for i in g) for g in groups)
+    # Contiguous split would give 200 vs 20; LPT lands at 110 vs 110
+    assert group_weights == [110.0, 110.0]
+    assert sorted(i for g in groups for i in g) == items
+
+
+def test_lpt_partition_more_groups_than_items() -> None:
+    """Surplus groups come back empty; every item still lands somewhere."""
+    items = ["a", "b"]
+    groups = helpers.lpt_partition(items, {"a": 1.0, "b": 1.0}, 4)
+    assert len(groups) == 4
+    assert sorted(i for g in groups for i in g) == items
+    assert sum(not g for g in groups) == 2
+
+
+def test_lpt_partition_tie_determinism() -> None:
+    """Equal weights assign in input order, so output is reproducible."""
+    items = [f"i{n}" for n in range(4)]
+    weights = dict.fromkeys(items, 1.0)
+    assert helpers.lpt_partition(items, weights, 2) == [["i0", "i2"], ["i1", "i3"]]

tests/script/test_update_integration_test_durations.py

@@ -0,0 +1,130 @@
+"""Unit tests for script/update_integration_test_durations.py."""
+
+import json
+from pathlib import Path
+import sys
+from unittest.mock import patch
+
+import pytest
+
+# Add the script directory to Python path so we can import the module
+script_dir = str((Path(__file__).parent / ".." / ".." / "script").resolve())
+sys.path.insert(0, script_dir)
+
+import helpers  # noqa: E402
+import update_integration_test_durations as uitd  # noqa: E402
+
+JUNIT_TEMPLATE = """<?xml version="1.0" encoding="utf-8"?>
+<testsuites><testsuite>{testcases}</testsuite></testsuites>
+"""
+
+KNOWN = {
+    "tests/integration/test_a.py",
+    "tests/integration/test_b.py",
+}
+
+
+def _write_junit(path: Path, testcases: str) -> None:
+    path.write_text(JUNIT_TEMPLATE.format(testcases=testcases), encoding="utf-8")
+
+
+def test_collect_durations_sums_per_file(tmp_path: Path) -> None:
+    """Testcases from the same module sum."""
+    _write_junit(
+        tmp_path / "a.xml",
+        '<testcase classname="tests.integration.test_a" name="t1" time="1.5"/>'
+        '<testcase classname="tests.integration.test_a" name="t2" time="2.0"/>'
+        '<testcase classname="tests.integration.test_b" name="t1" time="4.0"/>',
+    )
+    assert uitd.collect_durations(tmp_path, KNOWN) == {
+        "tests/integration/test_a.py": 3.5,
+        "tests/integration/test_b.py": 4.0,
+    }
+
+
+def test_collect_durations_class_based_testcase(tmp_path: Path) -> None:
+    """A class-based classname still maps to its module file."""
+    _write_junit(
+        tmp_path / "a.xml",
+        '<testcase classname="tests.integration.test_a.TestFoo" name="t" time="2.5"/>',
+    )
+    assert uitd.collect_durations(tmp_path, KNOWN) == {
+        "tests/integration/test_a.py": 2.5
+    }
+
+
+def test_collect_durations_unknown_module_skipped(
+    tmp_path: Path, capsys: pytest.CaptureFixture[str]
+) -> None:
+    """A classname that maps to no known file is skipped with a warning."""
+    _write_junit(
+        tmp_path / "a.xml",
+        '<testcase classname="tests.integration.test_gone" name="t" time="2.5"/>',
+    )
+    assert uitd.collect_durations(tmp_path, KNOWN) == {}
+    assert "test_gone" in capsys.readouterr().err
+
+
+def test_collect_durations_skips_skipped_testcases(tmp_path: Path) -> None:
+    """Skipped testcases do not record a bogus zero duration."""
+    _write_junit(
+        tmp_path / "a.xml",
+        '<testcase classname="tests.integration.test_a" name="t" time="0">'
+        "<skipped/></testcase>",
+    )
+    assert uitd.collect_durations(tmp_path, KNOWN) == {}
+
+
+def test_collect_durations_unexpected_classname_aborts(tmp_path: Path) -> None:
+    """A classname outside tests.integration means the junit layout changed."""
+    _write_junit(
+        tmp_path / "a.xml",
+        '<testcase classname="tests.unit_tests.test_x" name="t" time="9.0"/>',
+    )
+    with pytest.raises(SystemExit):
+        uitd.collect_durations(tmp_path, KNOWN)
+
+
+def test_collect_durations_empty_dir_aborts(tmp_path: Path) -> None:
+    """No junit XML at all is a hard error, not an empty recording."""
+    with pytest.raises(SystemExit):
+        uitd.collect_durations(tmp_path, KNOWN)
+
+
+def test_main_merges_partial_run(tmp_path: Path) -> None:
+    """A partial run merges over the previous data instead of truncating it."""
+    tests_dir = tmp_path / "tests" / "integration"
+    tests_dir.mkdir(parents=True)
+    for name in ("test_a", "test_b", "test_c"):
+        (tests_dir / f"{name}.py").write_text("", encoding="utf-8")
+    durations_file = tmp_path / helpers.INTEGRATION_TEST_DURATIONS_FILE
+    durations_file.write_text(
+        json.dumps(
+            {
+                "tests/integration/test_a.py": 5.0,
+                "tests/integration/test_b.py": 7.0,
+                "tests/integration/test_gone.py": 9.0,
+            }
+        ),
+        encoding="utf-8",
+    )
+    junit_dir = tmp_path / "junit"
+    junit_dir.mkdir()
+    _write_junit(
+        junit_dir / "a.xml",
+        '<testcase classname="tests.integration.test_a" name="t" time="6.0"/>',
+    )
+    with (
+        patch.object(helpers, "root_path", str(tmp_path)),
+        patch.object(uitd, "DURATIONS_FILE", durations_file),
+    ):
+        # 1 of 3 files covered: refused without --allow-partial
+        with patch.object(sys, "argv", ["uitd", str(junit_dir)]):
+            assert uitd.main() == uitd.EXIT_LOW_COVERAGE
+        with patch.object(sys, "argv", ["uitd", str(junit_dir), "--allow-partial"]):
+            assert uitd.main() == 0
+    # test_a updated, test_b kept, deleted test_gone dropped
+    assert json.loads(durations_file.read_text()) == {
+        "tests/integration/test_a.py": 6.0,
+        "tests/integration/test_b.py": 7.0,
+    }