| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from vllm.model_executor.model_loader.mtp_validation import ( |
| 7 | disable_mtp_completeness_check, |
| 8 | is_mtp_completeness_check_enabled, |
| 9 | ) |
| 10 | |
| 11 | |
| 12 | def test_disable_mtp_completeness_check_is_scoped(): |
| 13 | assert is_mtp_completeness_check_enabled() |
| 14 | |
| 15 | with pytest.raises(RuntimeError), disable_mtp_completeness_check(): |
| 16 | assert not is_mtp_completeness_check_enabled() |
| 17 | raise RuntimeError |
| 18 | |
| 19 | assert is_mtp_completeness_check_enabled() |