Files
sglang/test/registered/unit/test_no_bare_pytest_main.py
T

91 lines
2.8 KiB
Python

import ast
import pathlib
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3]
_SCAN_ROOTS = [_REPO_ROOT / "python", _REPO_ROOT / "test"]
class TestNoBarePytestMain(CustomTestCase):
def test_no_bare_pytest_main_in_repo(self):
offenders = []
for root in _SCAN_ROOTS:
if not root.exists():
continue
for path in root.rglob("*.py"):
violation = _find_bare_pytest_main(path)
if violation is not None:
offenders.append(violation)
self.assertFalse(
offenders,
msg=(
"Found bare `pytest.main(...)` in __main__ blocks (must be "
"wrapped in sys.exit(...) so failing tests propagate the exit "
"code to the CI runner):\n " + "\n ".join(offenders)
),
)
def _find_bare_pytest_main(path: pathlib.Path):
"""Return `<rel_path>:<lineno>` if `path` has a bare pytest.main(...) call
inside `if __name__ == "__main__":`, else None."""
try:
source = path.read_text(encoding="utf-8")
except (OSError, UnicodeDecodeError):
return None
try:
tree = ast.parse(source, filename=str(path))
except SyntaxError:
return None
for node in ast.walk(tree):
if not isinstance(node, ast.If):
continue
if not _is_main_guard(node.test):
continue
for stmt in node.body:
if _is_bare_pytest_main_call(stmt):
rel = path.relative_to(_REPO_ROOT)
return f"{rel}:{stmt.lineno}"
return None
def _is_main_guard(test: ast.expr) -> bool:
"""Match `__name__ == "__main__"` (either side)."""
if not isinstance(test, ast.Compare) or len(test.ops) != 1:
return False
if not isinstance(test.ops[0], ast.Eq):
return False
sides = [test.left, *test.comparators]
has_name = any(isinstance(s, ast.Name) and s.id == "__name__" for s in sides)
has_main = any(isinstance(s, ast.Constant) and s.value == "__main__" for s in sides)
return has_name and has_main
def _is_bare_pytest_main_call(stmt: ast.stmt) -> bool:
"""Match `pytest.main(...)` whose return value is discarded.
`sys.exit(pytest.main(...))` and `code = pytest.main(...)` are fine."""
if not isinstance(stmt, ast.Expr):
return False
call = stmt.value
if not isinstance(call, ast.Call):
return False
func = call.func
return (
isinstance(func, ast.Attribute)
and func.attr == "main"
and isinstance(func.value, ast.Name)
and func.value.id == "pytest"
)
if __name__ == "__main__":
unittest.main()