anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Misc] Cleanup: use consistent pytest-mock in unit tests (#2698)

Browse Source 
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed 2026-04-13 19:21:42 +08:00
1 parent 2c67c30550
commit e0cdbe9a5d
26 files changed
+1598 -1369
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py
+75 -56
View File
@@ -15,9 +15,10 @@ import importlib.util
import os
import sys
import types
from unittest.mock import MagicMock, patch
import pytest
import torch
from pytest_mock import MockerFixture
# Direct file import to avoid vllm_omni.__init__ patch dependencies.
_BASE = os.path.join(
@@ -41,28 +42,31 @@ def _load_module(name: str, filename: str):
return mod
def _build_mock_modules() -> dict[str, object]:
def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
"""Build the dict of modules to inject into sys.modules."""
platforms_mock = MagicMock()
platforms_mock = mocker.MagicMock()
platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False
logger_mock = MagicMock()
logger_mock.init_logger = lambda name: MagicMock()
logger_mock = mocker.MagicMock()
logger_mock.init_logger = lambda name: mocker.MagicMock()
vllm_config_mod = MagicMock()
vllm_config_mod.set_current_vllm_config = lambda cfg: MagicMock(__enter__=MagicMock(), __exit__=MagicMock())
vllm_config_mod = mocker.MagicMock()
vllm_config_mod.set_current_vllm_config = lambda cfg: mocker.MagicMock(
__enter__=mocker.MagicMock(),
__exit__=mocker.MagicMock(),
)
weight_utils_mock = MagicMock()
weight_utils_mock = mocker.MagicMock()
weight_utils_mock.default_weight_loader = lambda p, w: None
pkg = types.ModuleType("vllm_omni.model_executor.models.qwen3_tts")
pkg.__path__ = [os.path.abspath(_BASE)]
return {
"vllm_omni": MagicMock(),
"vllm_omni": mocker.MagicMock(),
"vllm_omni.platforms": platforms_mock,
"vllm.logger": logger_mock,
"vllm.config": MagicMock(),
"vllm.config": mocker.MagicMock(),
"vllm.config.vllm": vllm_config_mod,
"vllm.model_executor.model_loader.weight_utils": weight_utils_mock,
"vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"),
@@ -71,38 +75,47 @@ def _build_mock_modules() -> dict[str, object]:
}
def _load_target_classes():
def _load_target_classes(mocker: MockerFixture):
"""Load config and code predictor modules with mocked dependencies.
Uses patch.dict to ensure sys.modules is always restored, even on failure.
Uses mocker.patch.dict to ensure sys.modules is always restored, even on failure.
"""
mocks = _build_mock_modules()
with patch.dict(sys.modules, mocks):
config_mod = _load_module(
"vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts",
"configuration_qwen3_tts.py",
)
sys.modules["vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts"] = config_mod
mocks = _build_mock_modules(mocker)
mocker.patch.dict(sys.modules, mocks)
config_mod = _load_module(
"vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts",
"configuration_qwen3_tts.py",
)
sys.modules["vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts"] = config_mod
cp_mod = _load_module(
"vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm",
"qwen3_tts_code_predictor_vllm.py",
)
cp_mod = _load_module(
"vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm",
"qwen3_tts_code_predictor_vllm.py",
)
return config_mod, cp_mod
_config_mod, _cp_mod = _load_target_classes()
Qwen3TTSTalkerCodePredictorConfig = _config_mod.Qwen3TTSTalkerCodePredictorConfig
Qwen3TTSTalkerConfig = _config_mod.Qwen3TTSTalkerConfig
CodePredictorWrapper = _cp_mod.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
CodePredictorModel = _cp_mod.Qwen3TTSTalkerCodePredictorModelVLLM
@pytest.fixture
def loaded_target_classes(mocker: MockerFixture):
config_mod, cp_mod = _load_target_classes(mocker)
return (
config_mod.Qwen3TTSTalkerCodePredictorConfig,
config_mod.Qwen3TTSTalkerConfig,
cp_mod.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
cp_mod.Qwen3TTSTalkerCodePredictorModelVLLM,
)
def _make_tiny_config() -> tuple:
def _make_tiny_config(loaded_target_classes) -> tuple:
"""Create minimal configs for a tiny code predictor model."""
cp_config = Qwen3TTSTalkerCodePredictorConfig(
(
qwen3_tts_talker_code_predictor_config,
qwen3_tts_talker_config,
_,
_,
) = loaded_target_classes
cp_config = qwen3_tts_talker_code_predictor_config(
vocab_size=64,
hidden_size=32,
intermediate_size=64,
@@ -113,16 +126,16 @@ def _make_tiny_config() -> tuple:
num_code_groups=4,
rms_norm_eps=1e-6,
)
talker_config = Qwen3TTSTalkerConfig(
talker_config = qwen3_tts_talker_config(
hidden_size=32,
num_code_groups=4,
)
return cp_config, talker_config
def _make_vllm_config(max_num_seqs: int = 4) -> MagicMock:
def _make_vllm_config(mocker: MockerFixture, max_num_seqs: int = 4):
"""Create a mock VllmConfig with scheduler_config."""
vllm_config = MagicMock()
vllm_config = mocker.MagicMock()
vllm_config.scheduler_config.max_num_seqs = max_num_seqs
return vllm_config
@@ -130,12 +143,13 @@ def _make_vllm_config(max_num_seqs: int = 4) -> MagicMock:
class TestCodePredictorDtypeAlignment:
"""Test that code predictor buffers match model parameter dtype."""
def test_ensure_buffers_uses_given_dtype(self) -> None:
def test_ensure_buffers_uses_given_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None:
"""_ensure_buffers should create proj_buf with the given dtype."""
cp_config, talker_config = _make_tiny_config()
vllm_config = _make_vllm_config()
_, _, code_predictor_wrapper, _ = loaded_target_classes
cp_config, talker_config = _make_tiny_config(loaded_target_classes)
vllm_config = _make_vllm_config(mocker)
predictor = CodePredictorWrapper(
predictor = code_predictor_wrapper(
vllm_config=vllm_config,
config=cp_config,
talker_config=talker_config,
@@ -150,12 +164,13 @@ class TestCodePredictorDtypeAlignment:
predictor._ensure_buffers(torch.device("cpu"), torch.float32)
assert predictor._proj_buf.dtype == torch.float32
def test_warmup_aligns_buffer_to_model_params(self) -> None:
def test_warmup_aligns_buffer_to_model_params(self, mocker: MockerFixture, loaded_target_classes) -> None:
"""_warmup_buckets should align proj_buf dtype to model parameters."""
cp_config, talker_config = _make_tiny_config()
vllm_config = _make_vllm_config(max_num_seqs=2)
_, _, code_predictor_wrapper, _ = loaded_target_classes
cp_config, talker_config = _make_tiny_config(loaded_target_classes)
vllm_config = _make_vllm_config(mocker, max_num_seqs=2)
predictor = CodePredictorWrapper(
predictor = code_predictor_wrapper(
vllm_config=vllm_config,
config=cp_config,
talker_config=talker_config,
@@ -177,12 +192,13 @@ class TestCodePredictorDtypeAlignment:
assert predictor._proj_buf.dtype == torch.float16
def test_setup_compile_caches_model_dtype(self) -> None:
def test_setup_compile_caches_model_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None:
"""_setup_compile should cache model parameter dtype."""
cp_config, talker_config = _make_tiny_config()
vllm_config = _make_vllm_config(max_num_seqs=2)
_, _, code_predictor_wrapper, _ = loaded_target_classes
cp_config, talker_config = _make_tiny_config(loaded_target_classes)
vllm_config = _make_vllm_config(mocker, max_num_seqs=2)
predictor = CodePredictorWrapper(
predictor = code_predictor_wrapper(
vllm_config=vllm_config,
config=cp_config,
talker_config=talker_config,
@@ -193,12 +209,13 @@ class TestCodePredictorDtypeAlignment:
predictor._setup_compile()
assert predictor._model_dtype == torch.float16
def test_forward_with_mismatched_input_dtype(self) -> None:
def test_forward_with_mismatched_input_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None:
"""forward() should not crash when inputs are float32 but model is float16."""
cp_config, talker_config = _make_tiny_config()
vllm_config = _make_vllm_config(max_num_seqs=2)
_, _, code_predictor_wrapper, _ = loaded_target_classes
cp_config, talker_config = _make_tiny_config(loaded_target_classes)
vllm_config = _make_vllm_config(mocker, max_num_seqs=2)
predictor = CodePredictorWrapper(
predictor = code_predictor_wrapper(
vllm_config=vllm_config,
config=cp_config,
talker_config=talker_config,
@@ -231,10 +248,11 @@ class TestCodePredictorDtypeAlignment:
class TestCodePredictorModelDtype:
"""Test the inner model forward with different dtypes."""
def test_model_forward_float16(self) -> None:
def test_model_forward_float16(self, loaded_target_classes) -> None:
"""Inner model forward should work in float16."""
cp_config, _ = _make_tiny_config()
model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float16)
_, _, _, code_predictor_model = loaded_target_classes
cp_config, _ = _make_tiny_config(loaded_target_classes)
model = code_predictor_model(cp_config, talker_hidden_size=32).to(torch.float16)
bsz, seq_len = 1, 4
inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float16)
@@ -244,10 +262,11 @@ class TestCodePredictorModelDtype:
assert output.dtype == torch.float16
assert output.shape == (bsz, seq_len, 32)
def test_model_forward_float32(self) -> None:
def test_model_forward_float32(self, loaded_target_classes) -> None:
"""Inner model forward should work in float32."""
cp_config, _ = _make_tiny_config()
model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float32)
_, _, _, code_predictor_model = loaded_target_classes
cp_config, _ = _make_tiny_config(loaded_target_classes)
model = code_predictor_model(cp_config, talker_hidden_size=32).to(torch.float32)
bsz, seq_len = 1, 4
inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float32)