main
py 98 lines 3.17 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Sequence
4 from unittest.mock import MagicMock, patch
5
6 import pytest
7
8 from vllm.config import VllmConfig
9 from vllm.inputs import PromptType
10 from vllm.outputs import PoolingRequestOutput
11 from vllm.plugins.io_processors import get_io_processor
12 from vllm.plugins.io_processors.interface import IOProcessor
13 from vllm.renderers import BaseRenderer
14
15
16 class DummyIOProcessor(IOProcessor):
17 """Minimal IOProcessor used as the target of the mocked plugin entry point."""
18
19 def pre_process(
20 self,
21 prompt: object,
22 request_id: str | None = None,
23 **kwargs,
24 ) -> PromptType | Sequence[PromptType]:
25 raise NotImplementedError
26
27 def post_process(
28 self,
29 model_output: Sequence[PoolingRequestOutput],
30 request_id: str | None = None,
31 **kwargs,
32 ) -> object:
33 raise NotImplementedError
34
35
36 @pytest.fixture
37 def my_plugin_entry_points():
38 """Patch importlib.metadata.entry_points to expose a single 'my_plugin'
39 entry point backed by DummyIOProcessor, exercising the full plugin-loading
40 code path: entry_points → plugin.load() → func() →
41 resolve_obj_by_qualname → IOProcessor.__init__."""
42 qualname = f"{DummyIOProcessor.__module__}.{DummyIOProcessor.__qualname__}"
43 ep = MagicMock()
44 ep.name = "my_plugin"
45 ep.value = qualname
46 ep.load.return_value = lambda: qualname
47 with patch("importlib.metadata.entry_points", return_value=[ep]):
48 yield
49
50
51 def test_loading_missing_plugin():
52 vllm_config = VllmConfig()
53 renderer = MagicMock(spec=BaseRenderer)
54 with pytest.raises(ValueError):
55 get_io_processor(
56 vllm_config, renderer=renderer, plugin_from_init="wrong_plugin"
57 )
58
59
60 def test_loading_plugin(my_plugin_entry_points):
61 # Plugin name supplied via plugin_from_init.
62 vllm_config = MagicMock(spec=VllmConfig)
63 renderer = MagicMock(spec=BaseRenderer)
64
65 result = get_io_processor(
66 vllm_config, renderer=renderer, plugin_from_init="my_plugin"
67 )
68
69 assert isinstance(result, DummyIOProcessor)
70
71
72 def test_loading_missing_plugin_from_model_config():
73 # Build a mock VllmConfig whose hf_config advertises a plugin name,
74 # exercising the model-config code path without loading a real model.
75 mock_hf_config = MagicMock()
76 mock_hf_config.to_dict.return_value = {"io_processor_plugin": "wrong_plugin"}
77
78 vllm_config = MagicMock(spec=VllmConfig)
79 vllm_config.model_config.hf_config = mock_hf_config
80
81 renderer = MagicMock(spec=BaseRenderer)
82 with pytest.raises(ValueError):
83 get_io_processor(vllm_config, renderer=renderer)
84
85
86 def test_loading_plugin_from_model_config(my_plugin_entry_points):
87 # Plugin name supplied via the model's hf_config.
88 mock_hf_config = MagicMock()
89 mock_hf_config.to_dict.return_value = {"io_processor_plugin": "my_plugin"}
90
91 vllm_config = MagicMock(spec=VllmConfig)
92 vllm_config.model_config.hf_config = mock_hf_config
93
94 renderer = MagicMock(spec=BaseRenderer)
95
96 result = get_io_processor(vllm_config, renderer=renderer)
97
98 assert isinstance(result, DummyIOProcessor)