From 71b87793be45a1e16ba50d673937504f482247f3 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Thu, 2 Jul 2026 16:37:40 +0800 Subject: [PATCH] fix(ci): Execute newly added or modified scheduled_only tests in pre-submits --- .github/workflows/run_pathways_tests.yml | 6 +- .../workflows/run_tests_against_package.yml | 7 +- pytest.ini | 1 + tests/conftest.py | 11 + tests/unit/newly_added_detection_test.py | 194 +++++++++++++++++ tests/utils/newly_added_detection.py | 206 ++++++++++++++++++ 6 files changed, 423 insertions(+), 2 deletions(-) create mode 100644 tests/unit/newly_added_detection_test.py create mode 100644 tests/utils/newly_added_detection.py diff --git a/.github/workflows/run_pathways_tests.yml b/.github/workflows/run_pathways_tests.yml index 19adc362c4..ce03495541 100644 --- a/.github/workflows/run_pathways_tests.yml +++ b/.github/workflows/run_pathways_tests.yml @@ -113,7 +113,11 @@ jobs: if [ "${{ inputs.is_scheduled_run }}" = "true" ]; then FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}" else - FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and not scheduled_only" + if [ -z "${{ inputs.pytest_marker }}" ]; then + FINAL_PYTEST_MARKER="not scheduled_only or newly_added" + else + FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and (not scheduled_only or newly_added)" + fi fi export MAXTEXT_REPO_ROOT=$(pwd) export MAXTEXT_ASSETS_ROOT=$(pwd)/src/maxtext/assets diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 4b82a118e0..057a6329fa 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -100,6 +100,7 @@ jobs: uses: actions/checkout@v5 with: ref: ${{ inputs.maxtext_sha }} + fetch-depth: 0 - name: Download the maxtext wheel if: ${{ !inputs.maxtext_installed }} uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0 @@ -147,7 +148,11 @@ jobs: if [ "${INPUTS_IS_SCHEDULED_RUN}" == "true" ]; then FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER}" else - FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER} and not scheduled_only" + if [ -z "${INPUTS_PYTEST_MARKER}" ]; then + FINAL_PYTEST_MARKER="not scheduled_only or newly_added" + else + FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER} and (not scheduled_only or newly_added)" + fi fi # TODO: Use package data for testing and remove the env vars export MAXTEXT_REPO_ROOT=$(pwd) diff --git a/pytest.ini b/pytest.ini index 4976e05ebc..9cfe0e6103 100644 --- a/pytest.ini +++ b/pytest.ini @@ -38,6 +38,7 @@ markers = NOTE: this marker is not to be used manually, it is auto- applied to tests with external_* or tpu_only marker. scheduled_only: marks tests to run only on scheduled CI runs + newly_added: newly introduced or modified tests in PRs, executed even if scheduled_only integration_test: tests exercising larger portions of the system, including interactions with other systems like GCS, e.g., end_to_end tests diff --git a/tests/conftest.py b/tests/conftest.py index eec4afa225..cfb66c419c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -115,6 +115,7 @@ def _custom_iter(self): import maxtext # pylint: disable=unused-import from maxtext.common.gcloud_stub import is_decoupled +from tests.utils.newly_added_detection import get_changed_tests # Configure JAX to use unsafe_rbg PRNG implementation to match main scripts. if is_decoupled(): @@ -149,6 +150,15 @@ def pytest_collection_modifyitems(config, items): - Deselect tests marked as external_serving/training in decoupled mode. - Mark remaining tests with the `decoupled` marker when running decoupled. """ + + changed_tests = get_changed_tests() # set[(file_path, test_name)] + if changed_tests: + for item in items: + item_file = item.nodeid.split("::", 1)[0] + base_name = getattr(item, "originalname", item.name) + if (item_file, base_name) in changed_tests or (item_file, item.name) in changed_tests: + item.add_marker(pytest.mark.newly_added) + decoupled = is_decoupled() remaining = [] deselected = [] @@ -199,6 +209,7 @@ def pytest_configure(config): "external_training: goodput integrations", "decoupled: marked on tests that are not skipped due to GCP deps, when DECOUPLE_GCLOUD=TRUE", "skip_on_tpu7x: skip test if running on TPU7x platform", + "newly_added: newly introduced or modified tests in PRs, executed even if scheduled_only", ]: config.addinivalue_line("markers", m) diff --git a/tests/unit/newly_added_detection_test.py b/tests/unit/newly_added_detection_test.py new file mode 100644 index 0000000000..6c718a4d43 --- /dev/null +++ b/tests/unit/newly_added_detection_test.py @@ -0,0 +1,194 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the line-range test-change detector in ``newly_added_detection``.""" + +import textwrap + +import pytest + +from tests.utils.newly_added_detection import _build_diff_commands +from tests.utils.newly_added_detection import _is_test_file +from tests.utils.newly_added_detection import parse_changed_line_map +from tests.utils.newly_added_detection import touched_test_names + +pytestmark = pytest.mark.cpu_only + + +# --- _is_test_file ----------------------------------------------------------- + + +def test_is_test_file_accepts_test_suffixes_under_tests(): + assert _is_test_file("tests/unit/a_test.py") + assert _is_test_file("tests/integration/b_tests.py") + + +def test_is_test_file_rejects_non_tests_and_non_suffix(): + assert not _is_test_file("tests/utils/test_helpers.py") # helper, not *_test.py + assert not _is_test_file("src/maxtext/foo_test.py") # outside tests/ + assert not _is_test_file("tests/unit/helper.py") # not a test suffix + + +# --- _build_diff_commands ---------------------------------------------------- + + +def test_build_diff_commands_uses_only_merge_base_threedot_ranges(): + # Only three-dot (merge-base) ranges: remote first (CI / local-with-origin), + # then local (developer without an origin remote). Two-dot tip-vs-tip ranges are + # excluded because they over-report every commit the base gained past the fork + # point, dragging unrelated scheduled_only tests into pre-submit. + cmds = _build_diff_commands("main") + assert cmds == [ + ["git", "diff", "--unified=0", "origin/main...HEAD"], + ["git", "diff", "--unified=0", "main...HEAD"], + ] + # Guard against a two-dot form sneaking back in: every range arg is three-dot. + range_args = [cmd[-1] for cmd in cmds] + assert all("..." in arg for arg in range_args) + + +def test_build_diff_commands_honours_non_main_base(): + cmds = _build_diff_commands("release/v2") + assert cmds == [ + ["git", "diff", "--unified=0", "origin/release/v2...HEAD"], + ["git", "diff", "--unified=0", "release/v2...HEAD"], + ] + + +# --- parse_changed_line_map -------------------------------------------------- + + +def test_line_map_added_region_from_header_range(): + diff = ( + "diff --git a/tests/unit/a_test.py b/tests/unit/a_test.py\n" + "--- a/tests/unit/a_test.py\n" + "+++ b/tests/unit/a_test.py\n" + "@@ -0,0 +1,3 @@\n" + "+line1\n+line2\n+line3\n" + ) + assert parse_changed_line_map(diff) == {"tests/unit/a_test.py": {1, 2, 3}} + + +def test_line_map_single_line_default_length(): + diff = "+++ b/tests/unit/a_test.py\n" "@@ -10 +12 @@\n" "+changed\n" + assert parse_changed_line_map(diff) == {"tests/unit/a_test.py": {12}} + + +def test_line_map_pure_deletion_marks_boundary(): + # `+4,0` = pure deletion; the join sits between new-file lines 4 and 5. + diff = "+++ b/tests/unit/a_test.py\n" "@@ -5,2 +4,0 @@\n" "-gone1\n-gone2\n" + assert parse_changed_line_map(diff) == {"tests/unit/a_test.py": {4, 5}} + + +def test_line_map_accumulates_multiple_hunks_and_files(): + diff = ( + "+++ b/tests/unit/a_test.py\n" + "@@ -0,0 +1,1 @@\n" + "+x\n" + "@@ -8,0 +10,2 @@\n" + "+y\n+z\n" + "+++ b/tests/unit/b_test.py\n" + "@@ -0,0 +3,1 @@\n" + "+w\n" + ) + assert parse_changed_line_map(diff) == { + "tests/unit/a_test.py": {1, 10, 11}, + "tests/unit/b_test.py": {3}, + } + + +def test_line_map_ignores_deleted_file(): + diff = "+++ /dev/null\n" "@@ -1,2 +0,0 @@\n" "-a\n-b\n" + assert not parse_changed_line_map(diff) + + +# --- tests_touching_lines ---------------------------------------------------- + +_SOURCE = textwrap.dedent( + """\ + import pytest + + + class TestAlpha: + + @pytest.mark.scheduled_only + def test_existing(self): + x = 1 + assert x == 1 + + def test_untouched(self): + assert True + + + def test_top_level_untouched(): + assert True + + + async def test_async_new(): + assert True + """ +) +# Line numbers in _SOURCE: +# 6 @pytest.mark.scheduled_only +# 7 def test_existing (span 6-9, decorator included) +# 11 def test_untouched (span 11-12) +# 15 def test_top_level_untouched (span 15-16) +# 19 async def test_async_new (span 19-20) + + +def test_touching_body_line_flags_that_test(): + assert touched_test_names(_SOURCE, {9}) == {"test_existing"} + + +def test_touching_decorator_line_flags_that_test(): + assert touched_test_names(_SOURCE, {6}) == {"test_existing"} + + +def test_touching_async_test_is_detected(): + assert touched_test_names(_SOURCE, {20}) == {"test_async_new"} + + +def test_untouched_tests_are_not_flagged(): + # Editing test_existing must not drag in the neighbouring untouched tests. + assert touched_test_names(_SOURCE, {8, 9}) == {"test_existing"} + + +def test_lines_outside_any_test_flag_nothing(): + assert touched_test_names(_SOURCE, {1, 2}) == set() # imports / blank lines + + +def test_empty_touched_set_returns_empty(): + assert touched_test_names(_SOURCE, set()) == set() + + +def test_unparseable_source_returns_empty_gracefully(): + # pytest collection surfaces the syntax error itself; the parser must not raise. + assert touched_test_names("def broken(:\n", {1}) == set() + + +def test_insertion_after_untouched_test_does_not_flag_it(): + # The regression that git's hunk-header heuristic caused: a new test added + # right after an untouched test must flag ONLY the new test. + source = textwrap.dedent( + """\ + def test_old(): + assert True + + + def test_new(): + assert True + """ + ) + # test_old spans lines 1-2; test_new spans lines 5-6. The added lines are 5-6. + assert touched_test_names(source, {5, 6}) == {"test_new"} diff --git a/tests/utils/newly_added_detection.py b/tests/utils/newly_added_detection.py new file mode 100644 index 0000000000..314dcaaae4 --- /dev/null +++ b/tests/utils/newly_added_detection.py @@ -0,0 +1,206 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Detect the tests a PR adds or modifies, for scheduled_only pre-submit verification. + +The pre-submit CI runs ``pytest -m " and (not scheduled_only or newly_added)"``. +This module supplies the ``newly_added`` set: the tests a pull request touched, so that a +newly added or modified ``scheduled_only`` test runs at least once before merge instead of +being silently skipped until the nightly scheduled pipeline. + +Detection maps changed *line numbers* (from ``git diff --unified=0``) onto each test's +line span (from an ``ast`` parse of the new file). A test counts as changed only when a +changed line lands inside its own span. This avoids trusting git's hunk-header function +name, which points at the function *preceding* an insertion and would otherwise flag an +untouched test that merely sits above newly added code. + +Only the Python standard library is used, since this runs on a bare CI runner where +MaxText is not necessarily importable. +""" + +import ast +import os +import re +import subprocess + +# Matches pytest.ini's ``python_files = *_test.py *_tests.py``. +_TEST_SUFFIXES = ("_test.py", "_tests.py") +# Matches pytest.ini's ``testpaths = tests``. +_TESTS_ROOT = "tests/" +# Captures the new-file start line and (optional) length from a unified-diff hunk header: +# @@ -[,] +[,] @@ +_HUNK_RE = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") + + +def _is_test_file(path): + """Return True if ``path`` is a MaxText test file (under tests/ with a test suffix).""" + normalized = path.strip().replace(os.sep, "/") + return normalized.startswith(_TESTS_ROOT) and normalized.endswith(_TEST_SUFFIXES) + + +def parse_changed_line_map(diff_text): + """Map each changed file to the set of new-file line numbers it touched. + + Args: + diff_text: Raw ``git diff --unified=0`` output. + + Returns: + A dict of ``{file_path: set_of_new_line_numbers}``. ``file_path`` is repo-root + relative (the ``b/`` prefix is stripped). For a pure deletion (new length 0) the + two lines bracketing the removal are recorded, so a test that had lines removed is + still detected. Deleted files (``+++ /dev/null``) are omitted. + """ + line_map = {} + current_file = None + for line in diff_text.splitlines(): + if line.startswith("+++"): + path = line[3:].strip() + if path.startswith("b/"): + path = path[2:] + current_file = None if path == "/dev/null" else path + continue + if line.startswith("@@"): + if current_file is None: + continue + match = _HUNK_RE.match(line) + if match is None: + continue + new_start = int(match.group(1)) + new_len = int(match.group(2)) if match.group(2) is not None else 1 + touched = line_map.setdefault(current_file, set()) + if new_len > 0: + touched.update(range(new_start, new_start + new_len)) + else: + touched.update({new_start, new_start + 1}) + return line_map + + +def _iter_test_defs(tree): + """Yield ``(name, start_line, end_line)`` for every test function in an AST. + + Covers module-level test functions and methods declared directly inside a class, + which is what pytest collects. The span starts at the first decorator (if any) so + decorator-only edits are attributed to the test they decorate. + """ + def_types = (ast.FunctionDef, ast.AsyncFunctionDef) + for node in tree.body: + if isinstance(node, def_types) and node.name.startswith("test_"): + yield node.name, _span_start(node), node.end_lineno + elif isinstance(node, ast.ClassDef): + for sub in node.body: + if isinstance(sub, def_types) and sub.name.startswith("test_"): + yield sub.name, _span_start(sub), sub.end_lineno + + +def _span_start(node): + """Return the first source line of a def, including any decorator lines above it.""" + start = node.lineno + if node.decorator_list: + start = min(start, min(dec.lineno for dec in node.decorator_list)) + return start + + +def touched_test_names(source, touched_lines): + """Return the names of test functions in ``source`` whose span includes a changed line. + + Args: + source: The new file's Python source. + touched_lines: Set of changed new-file line numbers for that file. + + Returns: + A set of test function names. Empty if ``touched_lines`` is empty or ``source`` does + not parse (pytest collection surfaces a syntax error on its own, so raising here would + only hide it). + """ + if not touched_lines: + return set() + try: + tree = ast.parse(source) + except SyntaxError: + return set() + found = set() + for name, start, end in _iter_test_defs(tree): + if any(start <= line <= end for line in touched_lines): + found.add(name) + return found + + +def _build_diff_commands(base): + """Return the ordered ``git diff`` argument lists to try, most-precise first. + + Both ranges are three-dot (merge-base) ranges, so only the branch's own commits + are reported. Two-dot tip-vs-tip ranges are deliberately excluded: they over-report + every commit the base gained past the fork point (a stale local ``main`` can inflate + the changed set many-fold), which would drag unrelated ``scheduled_only`` tests into + pre-submit. The remote range covers CI and local checkouts that have an ``origin`` + remote; the local range is the fallback for a developer without ``origin`` and is + never reached in CI, where ``origin/`` is always fetched first. + """ + return [ + ["git", "diff", "--unified=0", f"origin/{base}...HEAD"], + ["git", "diff", "--unified=0", f"{base}...HEAD"], + ] + + +def get_changed_tests(base_ref=None): + """Return ``(file_path, test_name)`` for every test the PR added or modified. + + Args: + base_ref: Base git ref to diff against. Defaults to the ``GITHUB_BASE_REF`` + environment variable, then ``"main"``. + + Returns: + A set of ``(file_path, test_name)`` tuples, or an empty set when not in a git work + tree or when the diff cannot be computed. + """ + base = base_ref or os.environ.get("GITHUB_BASE_REF") or "main" + try: + inside = subprocess.run( + ["git", "rev-parse", "--is-inside-work-tree"], + capture_output=True, + check=False, + ) + if inside.returncode != 0: + return set() + if os.environ.get("GITHUB_ACTIONS") == "true": + subprocess.run( + ["git", "fetch", "origin", f"{base}:refs/remotes/origin/{base}"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + diff_text = None + for command in _build_diff_commands(base): + try: + diff_text = subprocess.check_output(command, text=True, stderr=subprocess.DEVNULL) + break + except Exception: # pylint: disable=broad-exception-caught + continue + if diff_text is None: + return set() + except Exception: # pylint: disable=broad-exception-caught + return set() + + changed = set() + for path, touched_lines in parse_changed_line_map(diff_text).items(): + if not _is_test_file(path): + continue + try: + with open(path, "r", encoding="utf-8") as handle: + source = handle.read() + except OSError: + continue + for name in touched_test_names(source, touched_lines): + changed.add((path, name)) + return changed