| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import contextmanager |
| 5 | from unittest.mock import MagicMock, patch |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.platforms.interface import DeviceCapability |
| 10 | |
| 11 | |
| 12 | @pytest.fixture |
| 13 | def mock_cuda_platform(): |
| 14 | """Fixture that returns a factory for creating mocked CUDA platforms. |
| 15 | |
| 16 | Usage: |
| 17 | def test_something(mock_cuda_platform): |
| 18 | with mock_cuda_platform(is_cuda=True, capability=(9, 0)): |
| 19 | # test code |
| 20 | """ |
| 21 | |
| 22 | @contextmanager |
| 23 | def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = None): |
| 24 | mock_platform = MagicMock() |
| 25 | mock_platform.is_cuda.return_value = is_cuda |
| 26 | mock_platform.is_xpu.return_value = False |
| 27 | device_capability = ( |
| 28 | DeviceCapability(*capability) if capability is not None else None |
| 29 | ) |
| 30 | mock_platform.get_device_capability.return_value = device_capability |
| 31 | |
| 32 | def is_device_capability_family( |
| 33 | requested_capability: int, device_id: int = 0 |
| 34 | ) -> bool: |
| 35 | current_capability = mock_platform.get_device_capability( |
| 36 | device_id=device_id |
| 37 | ) |
| 38 | if current_capability is None: |
| 39 | return False |
| 40 | return current_capability.major == (requested_capability // 10) |
| 41 | |
| 42 | mock_platform.is_device_capability_family.side_effect = ( |
| 43 | is_device_capability_family |
| 44 | ) |
| 45 | with patch("vllm.platforms.current_platform", mock_platform): |
| 46 | yield mock_platform |
| 47 | |
| 48 | return _mock_platform |
| 49 | |
| 50 | |
| 51 | @pytest.fixture |
| 52 | def mock_xpu_platform(): |
| 53 | """Fixture that returns a factory for creating mocked XPU platforms. |
| 54 | |
| 55 | Usage: |
| 56 | def test_something(mock_xpu_platform): |
| 57 | with mock_xpu_platform(): |
| 58 | # test code |
| 59 | """ |
| 60 | |
| 61 | @contextmanager |
| 62 | def _mock_platform(): |
| 63 | mock_platform = MagicMock() |
| 64 | mock_platform.is_cuda.return_value = False |
| 65 | mock_platform.is_xpu.return_value = True |
| 66 | with patch("vllm.platforms.current_platform", mock_platform): |
| 67 | yield mock_platform |
| 68 | |
| 69 | return _mock_platform |