diff --git a/tests/engine/test_engine_type_consistency.py b/tests/engine/test_engine_type_consistency.py new file mode 100644 index 0000000..63aa228 --- /dev/null +++ b/tests/engine/test_engine_type_consistency.py @@ -0,0 +1,41 @@ +"""Guard against drift between the C++ and Python EngineType enums. + +csrc/engine/scheduler.hpp and deepinfer/engine/engine_type.py define the same +enum twice (the C++ side is not exposed through pybind), and the two copies +are only kept in sync by a comment convention. Parse the header and assert +both definitions match, so an accidental drift fails fast in CI instead of a +Python enum value silently being interpreted as a different engine type by +the C++ scheduler. +""" + +import re +from pathlib import Path + +import pytest + +from deepinfer.engine.engine_type import EngineType + +pytestmark = [pytest.mark.ci, pytest.mark.ci_npu] + +HEADER = Path(__file__).resolve().parents[2] / "csrc" / "engine" / "scheduler.hpp" + +ENUM_RE = re.compile(r"enum\s+class\s+EngineType\s*:\s*int\s*\{(?P.*?)\};", re.DOTALL) +MEMBER_RE = re.compile(r"(\w+)\s*=\s*(-?\d+)") + + +def cpp_engine_types() -> dict[str, int]: + source = HEADER.read_text(encoding="utf-8") + match = ENUM_RE.search(source) + assert match is not None, f"enum class EngineType not found in {HEADER}" + members = MEMBER_RE.findall(match.group("body")) + assert members, f"no members parsed from enum EngineType in {HEADER}" + return {name: int(value) for name, value in members} + + +def test_engine_type_matches_cpp_definition() -> None: + cpp = cpp_engine_types() + py = {member.name: member.value for member in EngineType} + assert py == cpp, ( + "EngineType definitions drifted between csrc/engine/scheduler.hpp and " + f"deepinfer/engine/engine_type.py:\n C++: {cpp}\n Python: {py}" + )