diff --git a/pyproject.toml b/pyproject.toml index 965a3e4b2..9d9c0e933 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ dev = [ "pytest>=7.0.0", "pytest-asyncio>=0.21.0", "pytest-cov>=4.0.0", + "pytest-mock>=3.10.0", "mypy==1.11.1", "pre-commit==4.0.1", "openai-whisper>=20250625", diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index ea743eadd..39b7f84fb 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -7,9 +7,9 @@ Unit tests for patch.py import asyncio import json -from unittest.mock import AsyncMock, MagicMock import pytest +from pytest_mock import MockerFixture from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput, async_request_openai_chat_omni_completions @@ -51,7 +51,7 @@ def create_sse_chunk(data_dict): @pytest.mark.asyncio -async def test_output_tokens_assigned_with_metrics(): +async def test_output_tokens_assigned_with_metrics(mocker: MockerFixture): """Test that output.output_tokens is assigned when metrics are present""" # Arrange request_input = RequestFuncInput( @@ -82,8 +82,8 @@ async def test_output_tokens_assigned_with_metrics(): ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -95,7 +95,7 @@ async def test_output_tokens_assigned_with_metrics(): @pytest.mark.asyncio -async def test_output_tokens_not_assigned_without_metrics(): +async def test_output_tokens_not_assigned_without_metrics(mocker: MockerFixture): """Test that output.output_tokens defaults to 0 when no metrics present""" # Arrange request_input = RequestFuncInput( @@ -125,8 +125,8 @@ async def test_output_tokens_not_assigned_without_metrics(): ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -138,7 +138,7 @@ async def test_output_tokens_not_assigned_without_metrics(): @pytest.mark.asyncio -async def test_output_tokens_assigned_multiple_metrics(): +async def test_output_tokens_assigned_multiple_metrics(mocker: MockerFixture): """Test that output.output_tokens is updated with the latest metrics value""" # Arrange request_input = RequestFuncInput( @@ -177,8 +177,8 @@ async def test_output_tokens_assigned_multiple_metrics(): ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -190,7 +190,7 @@ async def test_output_tokens_assigned_multiple_metrics(): @pytest.mark.asyncio -async def test_output_tokens_with_audio_and_text(): +async def test_output_tokens_with_audio_and_text(mocker: MockerFixture): """Test output_tokens assignment in mixed audio and text response""" # Arrange request_input = RequestFuncInput( @@ -226,8 +226,8 @@ async def test_output_tokens_with_audio_and_text(): ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -238,7 +238,7 @@ async def test_output_tokens_with_audio_and_text(): @pytest.mark.asyncio -async def test_output_tokens_with_missing_num_tokens_out(): +async def test_output_tokens_with_missing_num_tokens_out(mocker: MockerFixture): """Test that output_tokens defaults to 0 when num_tokens_out is missing in metrics""" # Arrange request_input = RequestFuncInput( @@ -269,8 +269,8 @@ async def test_output_tokens_with_missing_num_tokens_out(): ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -312,7 +312,7 @@ class TestTextLatencyAttribute: assert text_latency is not None or hasattr(output, "text_latency"), "text_latency attribute should exist" @pytest.mark.asyncio - async def test_text_latency_assigned_with_text_response(self): + async def test_text_latency_assigned_with_text_response(self, mocker: MockerFixture): """Test that text_latency is assigned when text response is received""" request_input = RequestFuncInput( model="test-model", @@ -341,8 +341,8 @@ class TestTextLatencyAttribute: ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -353,7 +353,7 @@ class TestTextLatencyAttribute: assert output.text_latency > 0, "text_latency should be greater than 0 for text response" @pytest.mark.asyncio - async def test_text_latency_updated_with_multiple_text_chunks(self): + async def test_text_latency_updated_with_multiple_text_chunks(self, mocker: MockerFixture): """Test that text_latency is updated as more text chunks arrive""" request_input = RequestFuncInput( model="test-model", @@ -388,8 +388,8 @@ class TestTextLatencyAttribute: ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -401,7 +401,7 @@ class TestTextLatencyAttribute: assert output.generated_text == "First second third" @pytest.mark.asyncio - async def test_text_latency_with_only_audio_response(self): + async def test_text_latency_with_only_audio_response(self, mocker: MockerFixture): """Test text_latency behavior when only audio is received""" request_input = RequestFuncInput( model="test-model", @@ -424,8 +424,8 @@ class TestTextLatencyAttribute: ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -437,7 +437,7 @@ class TestTextLatencyAttribute: assert output.text_latency >= 0, "text_latency should be non-negative" @pytest.mark.asyncio - async def test_text_latency_not_affected_by_metrics(self): + async def test_text_latency_not_affected_by_metrics(self, mocker: MockerFixture): """Test that text_latency is independent of metrics data""" request_input = RequestFuncInput( model="test-model", @@ -467,8 +467,8 @@ class TestTextLatencyAttribute: ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -481,7 +481,7 @@ class TestTextLatencyAttribute: assert output.output_tokens == 100, "metrics should not affect text_latency" @pytest.mark.asyncio - async def test_text_latency_mixed_modalities(self): + async def test_text_latency_mixed_modalities(self, mocker: MockerFixture): """Test text_latency with mixed text and audio modalities""" request_input = RequestFuncInput( model="test-model", @@ -516,8 +516,8 @@ class TestTextLatencyAttribute: ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) @@ -529,7 +529,7 @@ class TestTextLatencyAttribute: assert output.generated_text == "Text more text" @pytest.mark.asyncio - async def test_text_latency_value_consistency(self): + async def test_text_latency_value_consistency(self, mocker: MockerFixture): """Test that text_latency matches latency minus ttft relationship""" request_input = RequestFuncInput( model="test-model", @@ -557,8 +557,8 @@ class TestTextLatencyAttribute: ] mock_response = MockResponse(200, chunks) - mock_session = AsyncMock() - mock_session.post = MagicMock(return_value=mock_response) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) # Act output = await async_request_openai_chat_omni_completions(request_input, mock_session) diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 73cfa104f..e2bd7ef8a 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -10,10 +10,9 @@ This module tests the DiffusionWorker implementation: - wake_up: Waking worker from sleep mode """ -from unittest.mock import Mock, patch - import pytest import torch +from pytest_mock import MockerFixture from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker @@ -21,9 +20,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @pytest.fixture -def mock_od_config(): +def mock_od_config(mocker: MockerFixture): """Create a mock OmniDiffusionConfig.""" - config = Mock() + config = mocker.Mock() config.num_gpus = 1 config.master_port = 12345 config.enable_sleep_mode = False @@ -34,22 +33,23 @@ def mock_od_config(): @pytest.fixture -def mock_gpu_worker(mock_od_config): +def mock_gpu_worker(mocker: MockerFixture, mock_od_config): """Create a DiffusionWorker with mocked initialization.""" - with patch.object(DiffusionWorker, "init_device"), patch.object(DiffusionWorker, "load_model"): - worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config) - # Mock the model_runner with pipeline - worker.model_runner = Mock() - worker.model_runner.pipeline = Mock() - worker.device = torch.device("cuda", 0) - worker._sleep_saved_buffers = {} - return worker + mocker.patch.object(DiffusionWorker, "init_device") + mocker.patch.object(DiffusionWorker, "load_model") + worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config) + # Mock the model_runner with pipeline + worker.model_runner = mocker.Mock() + worker.model_runner.pipeline = mocker.Mock() + worker.device = torch.device("cuda", 0) + worker._sleep_saved_buffers = {} + return worker class TestDiffusionWorkerLoadWeights: """Test DiffusionWorker.load_weights method.""" - def test_load_weights_calls_pipeline(self, mock_gpu_worker): + def test_load_weights_calls_pipeline(self, mocker: MockerFixture, mock_gpu_worker): """Test that load_weights delegates to model_runner.load_weights.""" # Setup mock weights mock_weights = [ @@ -59,7 +59,7 @@ class TestDiffusionWorkerLoadWeights: expected_loaded = {"layer1.weight", "layer2.weight"} # Configure model_runner mock - mock_gpu_worker.model_runner.load_weights = Mock(return_value=expected_loaded) + mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=expected_loaded) # Call load_weights result = mock_gpu_worker.load_weights(mock_weights) @@ -68,9 +68,9 @@ class TestDiffusionWorkerLoadWeights: mock_gpu_worker.model_runner.load_weights.assert_called_once_with(mock_weights) assert result == expected_loaded - def test_load_weights_empty_iterable(self, mock_gpu_worker): + def test_load_weights_empty_iterable(self, mocker: MockerFixture, mock_gpu_worker): """Test load_weights with empty weights iterable.""" - mock_gpu_worker.model_runner.load_weights = Mock(return_value=set()) + mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=set()) result = mock_gpu_worker.load_weights([]) @@ -81,11 +81,12 @@ class TestDiffusionWorkerLoadWeights: class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" - @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") - @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") - @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker): + def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" + mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") + # Setup process-scoped memory mocks # Before sleep: 3GB used # After sleep: 1GB used (freed 2GB) @@ -95,9 +96,9 @@ class TestDiffusionWorkerSleep: ] # Setup allocator mock - mock_allocator = Mock() - mock_allocator_class.get_instance = Mock(return_value=mock_allocator) - mock_allocator.sleep = Mock() + mock_allocator = mocker.Mock() + mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) + mock_allocator.sleep = mocker.Mock() # Call sleep with level 1 result = mock_gpu_worker.sleep(level=1) @@ -108,11 +109,12 @@ class TestDiffusionWorkerSleep: # Verify buffers were NOT saved (level 1 doesn't save buffers) assert len(mock_gpu_worker._sleep_saved_buffers) == 0 - @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") - @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") - @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker): + def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" + mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") + # Setup process-scoped memory mocks mock_get_process_memory.side_effect = [ 5 * 1024**3, # Before sleep @@ -120,14 +122,14 @@ class TestDiffusionWorkerSleep: ] # Setup allocator mock - mock_allocator = Mock() - mock_allocator_class.get_instance = Mock(return_value=mock_allocator) - mock_allocator.sleep = Mock() + mock_allocator = mocker.Mock() + mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) + mock_allocator.sleep = mocker.Mock() # Mock pipeline buffers mock_buffer1 = torch.randn(10, 10) mock_buffer2 = torch.randn(20, 20) - mock_gpu_worker.model_runner.pipeline.named_buffers = Mock( + mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock( return_value=[ ("buffer1", mock_buffer1), ("buffer2", mock_buffer2), @@ -146,42 +148,32 @@ class TestDiffusionWorkerSleep: assert "buffer1" in mock_gpu_worker._sleep_saved_buffers assert "buffer2" in mock_gpu_worker._sleep_saved_buffers - @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") - @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") - @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_memory_freed_validation( - self, - mock_allocator_class, - mock_platform, - mock_get_process_memory, - mock_gpu_worker, - ): + def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker): """Test that sleep validates memory was actually freed.""" + mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") + # Simulate process memory increase (should trigger assertion error) mock_get_process_memory.side_effect = [ 1 * 1024**3, # Before sleep: 1GB used 3 * 1024**3, # After sleep: 3GB used (negative freed) ] - mock_allocator = Mock() - mock_allocator_class.get_instance = Mock(return_value=mock_allocator) - mock_allocator.sleep = Mock() + mock_allocator = mocker.Mock() + mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) + mock_allocator.sleep = mocker.Mock() # This should raise an assertion error with pytest.raises(AssertionError, match="Memory usage increased after sleeping"): mock_gpu_worker.sleep(level=1) - @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") - @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") - @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_falls_back_to_device_memory_when_nvml_unavailable( - self, - mock_allocator_class, - mock_platform, - mock_get_process_memory, - mock_gpu_worker, - ): + def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep uses device-scoped fallback when NVML is unavailable.""" + + mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") mock_get_process_memory.side_effect = [None, None] mock_platform.get_free_memory.side_effect = [ 1 * 1024**3, # Before sleep @@ -189,9 +181,9 @@ class TestDiffusionWorkerSleep: ] mock_platform.get_device_total_memory.return_value = 8 * 1024**3 - mock_allocator = Mock() - mock_allocator_class.get_instance = Mock(return_value=mock_allocator) - mock_allocator.sleep = Mock() + mock_allocator = mocker.Mock() + mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) + mock_allocator.sleep = mocker.Mock() result = mock_gpu_worker.sleep(level=1) @@ -202,13 +194,14 @@ class TestDiffusionWorkerSleep: class TestDiffusionWorkerWakeUp: """Test DiffusionWorker.wake_up method.""" - @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker): + def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up without saved buffers (level 1 sleep).""" + mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + # Setup allocator mock - mock_allocator = Mock() - mock_allocator_class.get_instance = Mock(return_value=mock_allocator) - mock_allocator.wake_up = Mock() + mock_allocator = mocker.Mock() + mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) + mock_allocator.wake_up = mocker.Mock() # Ensure no saved buffers mock_gpu_worker._sleep_saved_buffers = {} @@ -220,13 +213,14 @@ class TestDiffusionWorkerWakeUp: mock_allocator.wake_up.assert_called_once_with(["weights"]) assert result is True - @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_wake_up_with_buffers(self, mock_allocator_class, mock_gpu_worker): + def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up with saved buffers (level 2 sleep).""" + mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + # Setup allocator mock - mock_allocator = Mock() - mock_allocator_class.get_instance = Mock(return_value=mock_allocator) - mock_allocator.wake_up = Mock() + mock_allocator = mocker.Mock() + mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) + mock_allocator.wake_up = mocker.Mock() # Create saved buffers saved_buffer1 = torch.randn(10, 10) @@ -237,12 +231,12 @@ class TestDiffusionWorkerWakeUp: } # Mock pipeline buffers (these will be restored) - mock_buffer1 = Mock() - mock_buffer1.data = Mock() - mock_buffer2 = Mock() - mock_buffer2.data = Mock() + mock_buffer1 = mocker.Mock() + mock_buffer1.data = mocker.Mock() + mock_buffer2 = mocker.Mock() + mock_buffer2.data = mocker.Mock() - mock_gpu_worker.model_runner.pipeline.named_buffers = Mock( + mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock( return_value=[ ("buffer1", mock_buffer1), ("buffer2", mock_buffer2), @@ -263,13 +257,14 @@ class TestDiffusionWorkerWakeUp: assert len(mock_gpu_worker._sleep_saved_buffers) == 0 assert result is True - @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_wake_up_partial_buffer_restore(self, mock_allocator_class, mock_gpu_worker): + def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up only restores buffers that were saved.""" + mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + # Setup allocator mock - mock_allocator = Mock() - mock_allocator_class.get_instance = Mock(return_value=mock_allocator) - mock_allocator.wake_up = Mock() + mock_allocator = mocker.Mock() + mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) + mock_allocator.wake_up = mocker.Mock() # Only save buffer1, not buffer2 saved_buffer1 = torch.randn(10, 10) @@ -278,12 +273,12 @@ class TestDiffusionWorkerWakeUp: } # Mock pipeline has both buffers - mock_buffer1 = Mock() - mock_buffer1.data = Mock() - mock_buffer2 = Mock() - mock_buffer2.data = Mock() + mock_buffer1 = mocker.Mock() + mock_buffer1.data = mocker.Mock() + mock_buffer2 = mocker.Mock() + mock_buffer2.data = mocker.Mock() - mock_gpu_worker.model_runner.pipeline.named_buffers = Mock( + mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock( return_value=[ ("buffer1", mock_buffer1), ("buffer2", mock_buffer2), diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index 2023a89e1..fd0be6c98 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -13,9 +13,9 @@ This module tests the WorkerWrapperBase implementation: """ from typing import Any -from unittest.mock import Mock, patch import pytest +from pytest_mock import MockerFixture from vllm_omni.diffusion.worker.diffusion_worker import ( CustomPipelineWorkerExtension, @@ -23,15 +23,18 @@ from vllm_omni.diffusion.worker.diffusion_worker import ( WorkerWrapperBase, ) +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + # ------------------------------------------------------------------------- # Fixtures # ------------------------------------------------------------------------- @pytest.fixture -def mock_od_config(): +def mock_od_config(mocker: MockerFixture): """Create a mock OmniDiffusionConfig for use in tests.""" - config = Mock() + config = mocker.Mock() config.num_gpus = 1 config.master_port = 12345 config.enable_sleep_mode = False @@ -72,9 +75,10 @@ class MockCustomPipeline: class TestWorkerWrapperBaseInitialization: """Test WorkerWrapperBase initialization behavior.""" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_basic_initialization(self, mock_worker_init, mock_od_config): + def test_basic_initialization(self, mocker: MockerFixture, mock_od_config): """Test basic initialization without extensions.""" + mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, @@ -103,49 +107,49 @@ class TestWorkerWrapperBaseInitialization: class TestWorkerWrapperBaseExtension: """Test WorkerWrapperBase worker extension functionality.""" - def test_prepare_worker_class_without_extension(self, mock_od_config): + def test_prepare_worker_class_without_extension(self, mocker: MockerFixture, mock_od_config): """Test _prepare_worker_class without a worker extension.""" - with patch.object(DiffusionWorker, "__init__", return_value=None): - wrapper = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - ) - worker_class = wrapper._prepare_worker_class() - assert worker_class == DiffusionWorker + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + ) + worker_class = wrapper._prepare_worker_class() + assert worker_class == DiffusionWorker - def test_prepare_worker_class_with_extension_class(self, mock_od_config): + def test_prepare_worker_class_with_extension_class(self, mocker: MockerFixture, mock_od_config): """Test _prepare_worker_class with an explicit extension class.""" class TestExtension: def custom_method(self): return "extension_method" - with patch.object(DiffusionWorker, "__init__", return_value=None): - wrapper = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - worker_extension_cls=TestExtension, - ) + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=TestExtension, + ) - assert hasattr(wrapper.worker.__class__, "custom_method") - assert TestExtension in wrapper.worker.__class__.__bases__ + assert hasattr(wrapper.worker.__class__, "custom_method") + assert TestExtension in wrapper.worker.__class__.__bases__ - @patch("vllm.utils.import_utils.resolve_obj_by_qualname") - def test_prepare_worker_class_with_extension_string(self, mock_resolve, mock_od_config): + def test_prepare_worker_class_with_extension_string(self, mocker: MockerFixture, mock_od_config): """Test _prepare_worker_class with worker extension as string.""" + mock_resolve = mocker.patch("vllm.utils.import_utils.resolve_obj_by_qualname") mock_resolve.return_value = TestExtension - with patch.object(DiffusionWorker, "__init__", return_value=None): - wrapper = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension", - ) + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension", + ) - assert hasattr(wrapper.worker.__class__, "custom_method") + assert hasattr(wrapper.worker.__class__, "custom_method") # ------------------------------------------------------------------------- @@ -156,70 +160,70 @@ class TestWorkerWrapperBaseExtension: class TestWorkerWrapperBaseDelegation: """Test WorkerWrapperBase delegation to wrapped worker.""" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_generate_delegation(self, mock_worker_init, mock_od_config): + def test_generate_delegation(self, mocker: MockerFixture, mock_od_config): """Test that generate() delegates to worker.generate().""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - mock_output = Mock() - wrapper.worker.generate = Mock(return_value=mock_output) + mock_output = mocker.Mock() + wrapper.worker.generate = mocker.Mock(return_value=mock_output) - mock_requests = [Mock()] + mock_requests = [mocker.Mock()] result = wrapper.generate(mock_requests) wrapper.worker.generate.assert_called_once_with(mock_requests) assert result == mock_output - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_execute_model_delegation(self, mock_worker_init, mock_od_config): + def test_execute_model_delegation(self, mocker: MockerFixture, mock_od_config): """Test that execute_model() delegates to worker.execute_model().""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - mock_output = Mock() - wrapper.worker.execute_model = Mock(return_value=mock_output) + mock_output = mocker.Mock() + wrapper.worker.execute_model = mocker.Mock(return_value=mock_output) - mock_reqs = [Mock()] + mock_reqs = [mocker.Mock()] result = wrapper.execute_model(mock_reqs, mock_od_config) wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config) assert result == mock_output - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_load_weights_delegation(self, mock_worker_init, mock_od_config): + def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config): """Test that load_weights() delegates to worker.load_weights().""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) expected_result = {"weight1", "weight2"} - wrapper.worker.load_weights = Mock(return_value=expected_result) + wrapper.worker.load_weights = mocker.Mock(return_value=expected_result) - mock_weights = [("weight1", Mock()), ("weight2", Mock())] + mock_weights = [("weight1", mocker.Mock()), ("weight2", mocker.Mock())] result = wrapper.load_weights(mock_weights) wrapper.worker.load_weights.assert_called_once_with(mock_weights) assert result == expected_result - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_sleep_delegation(self, mock_worker_init, mock_od_config): + def test_sleep_delegation(self, mocker: MockerFixture, mock_od_config): """Test that sleep() delegates to worker.sleep().""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - wrapper.worker.sleep = Mock(return_value=True) + wrapper.worker.sleep = mocker.Mock(return_value=True) result = wrapper.sleep(level=1) wrapper.worker.sleep.assert_called_once_with(1) assert result is True - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_wake_up_delegation(self, mock_worker_init, mock_od_config): + def test_wake_up_delegation(self, mocker: MockerFixture, mock_od_config): """Test that wake_up() delegates to worker.wake_up().""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - wrapper.worker.wake_up = Mock(return_value=True) + wrapper.worker.wake_up = mocker.Mock(return_value=True) result = wrapper.wake_up(tags=["weights"]) wrapper.worker.wake_up.assert_called_once_with(["weights"]) assert result is True - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_shutdown_delegation(self, mock_worker_init, mock_od_config): + def test_shutdown_delegation(self, mocker: MockerFixture, mock_od_config): """Test that shutdown() delegates to worker.shutdown().""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - wrapper.worker.shutdown = Mock(return_value=None) + wrapper.worker.shutdown = mocker.Mock(return_value=None) result = wrapper.shutdown() wrapper.worker.shutdown.assert_called_once() @@ -234,39 +238,39 @@ class TestWorkerWrapperBaseDelegation: class TestWorkerWrapperBaseExecuteMethod: """Test WorkerWrapperBase.execute_method functionality.""" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_execute_method_success(self, mock_worker_init, mock_od_config): + def test_execute_method_success(self, mocker: MockerFixture, mock_od_config): """Test execute_method successfully calls worker method.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - wrapper.worker.test_method = Mock(return_value="method_result") + wrapper.worker.test_method = mocker.Mock(return_value="method_result") result = wrapper.execute_method("test_method", "arg1", kwarg1="value1") wrapper.worker.test_method.assert_called_once_with("arg1", kwarg1="value1") assert result == "method_result" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_execute_method_with_no_args(self, mock_worker_init, mock_od_config): + def test_execute_method_with_no_args(self, mocker: MockerFixture, mock_od_config): """Test execute_method with no arguments.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - wrapper.worker.no_args_method = Mock(return_value="no_args_result") + wrapper.worker.no_args_method = mocker.Mock(return_value="no_args_result") result = wrapper.execute_method("no_args_method") wrapper.worker.no_args_method.assert_called_once_with() assert result == "no_args_result" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_execute_method_error(self, mock_worker_init, mock_od_config): + def test_execute_method_error(self, mocker: MockerFixture, mock_od_config): """Test execute_method raises exception on error.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - wrapper.worker.error_method = Mock(side_effect=RuntimeError("Test error")) + wrapper.worker.error_method = mocker.Mock(side_effect=RuntimeError("Test error")) with pytest.raises(RuntimeError, match="Test error"): wrapper.execute_method("error_method") - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_execute_method_invalid_type(self, mock_worker_init, mock_od_config): + def test_execute_method_invalid_type(self, mocker: MockerFixture, mock_od_config): """Test execute_method with invalid method type.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) with pytest.raises(AssertionError, match="Method must be str"): @@ -281,26 +285,26 @@ class TestWorkerWrapperBaseExecuteMethod: class TestWorkerWrapperBaseGetAttr: """Test WorkerWrapperBase.__getattr__ delegation.""" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_getattr_delegation(self, mock_worker_init, mock_od_config): + def test_getattr_delegation(self, mocker: MockerFixture, mock_od_config): """Test __getattr__ delegates to worker attributes.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.custom_attribute = "test_value" assert wrapper.custom_attribute == "test_value" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_getattr_method_access(self, mock_worker_init, mock_od_config): + def test_getattr_method_access(self, mocker: MockerFixture, mock_od_config): """Test __getattr__ delegates to worker methods.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) - wrapper.worker.custom_method = Mock(return_value="method_result") + wrapper.worker.custom_method = mocker.Mock(return_value="method_result") result = wrapper.custom_method() wrapper.worker.custom_method.assert_called_once() assert result == "method_result" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_getattr_missing_attribute(self, mock_worker_init, mock_od_config): + def test_getattr_missing_attribute(self, mocker: MockerFixture, mock_od_config): """Test __getattr__ raises AttributeError for missing attributes.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) with pytest.raises(AttributeError): _ = wrapper.nonexistent_attribute @@ -314,47 +318,46 @@ class TestWorkerWrapperBaseGetAttr: class TestWorkerWrapperBaseEdgeCases: """Test WorkerWrapperBase edge cases and special scenarios.""" - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_extension_conflict_warning(self, mock_worker_init, mock_od_config, caplog): + def test_extension_conflict_warning(self, mocker: MockerFixture, mock_od_config, caplog): """Test a warning is logged when an extension conflicts with worker.""" + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) class ConflictExtension: def load_model(self): return "extension_load_model" - with patch.object(DiffusionWorker, "load_model"): - wrapper = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - worker_extension_cls=ConflictExtension, - ) - assert wrapper.worker is not None + mocker.patch.object(DiffusionWorker, "load_model") + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=ConflictExtension, + ) + assert wrapper.worker is not None - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_multiple_extensions_same_class(self, mock_worker_init, mock_od_config): + def test_multiple_extensions_same_class(self, mocker: MockerFixture, mock_od_config): """Test that applying same extension twice doesn't duplicate it.""" class TestExtension: def custom_method(self): return "extension" - with patch.object(DiffusionWorker, "__init__", return_value=None): - wrapper1 = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - worker_extension_cls=TestExtension, - ) - wrapper2 = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - worker_extension_cls=TestExtension, - ) + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper1 = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=TestExtension, + ) + wrapper2 = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=TestExtension, + ) - assert hasattr(wrapper1.worker, "custom_method") - assert hasattr(wrapper2.worker, "custom_method") + assert hasattr(wrapper1.worker, "custom_method") + assert hasattr(wrapper2.worker, "custom_method") # ------------------------------------------------------------------------- @@ -365,11 +368,12 @@ class TestWorkerWrapperBaseEdgeCases: class TestCustomPipelineWorkerExtension: """Test CustomPipelineWorkerExtension functionality.""" - @patch("torch.cuda.empty_cache") - @patch("gc.collect") - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_re_init_pipeline_basic(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + def test_re_init_pipeline_basic(self, mocker: MockerFixture, mock_od_config): """Test basic re_init_pipeline functionality.""" + mocker.patch("torch.cuda.empty_cache") + mocker.patch("gc.collect") + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, @@ -378,12 +382,12 @@ class TestCustomPipelineWorkerExtension: ) # Setup mock model_runner and pipeline - mock_model_runner = Mock() - mock_pipeline = Mock() + mock_model_runner = mocker.Mock() + mock_pipeline = mocker.Mock() mock_model_runner.pipeline = mock_pipeline wrapper.worker.model_runner = mock_model_runner - wrapper.worker.init_lora_manager = Mock() - wrapper.worker.load_model = Mock() + wrapper.worker.init_lora_manager = mocker.Mock() + wrapper.worker.load_model = mocker.Mock() custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} @@ -397,11 +401,12 @@ class TestCustomPipelineWorkerExtension: ) wrapper.worker.init_lora_manager.assert_called_once() - @patch("torch.cuda.empty_cache") - @patch("gc.collect") - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_re_init_pipeline_cleanup(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + def test_re_init_pipeline_cleanup(self, mocker: MockerFixture, mock_od_config): """Test that re_init_pipeline properly cleans up old pipeline.""" + mock_gc_collect = mocker.patch("gc.collect") + mock_empty_cache = mocker.patch("torch.cuda.empty_cache") + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, @@ -410,12 +415,12 @@ class TestCustomPipelineWorkerExtension: ) # Setup mock model_runner with pipeline - mock_model_runner = Mock() - mock_pipeline = Mock() + mock_model_runner = mocker.Mock() + mock_pipeline = mocker.Mock() mock_model_runner.pipeline = mock_pipeline wrapper.worker.model_runner = mock_model_runner - wrapper.worker.init_lora_manager = Mock() - wrapper.worker.load_model = Mock() + wrapper.worker.init_lora_manager = mocker.Mock() + wrapper.worker.load_model = mocker.Mock() custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} @@ -426,11 +431,12 @@ class TestCustomPipelineWorkerExtension: mock_gc_collect.assert_called_once() mock_empty_cache.assert_called_once() - @patch("torch.cuda.empty_cache") - @patch("gc.collect") - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_re_init_pipeline_none_pipeline(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + def test_re_init_pipeline_none_pipeline(self, mocker: MockerFixture, mock_od_config): """Test re_init_pipeline when pipeline is None.""" + mocker.patch("torch.cuda.empty_cache") + mocker.patch("gc.collect") + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, @@ -439,11 +445,11 @@ class TestCustomPipelineWorkerExtension: ) # Setup mock model_runner with None pipeline - mock_model_runner = Mock() + mock_model_runner = mocker.Mock() mock_model_runner.pipeline = None wrapper.worker.model_runner = mock_model_runner - wrapper.worker.init_lora_manager = Mock() - wrapper.worker.load_model = Mock() + wrapper.worker.init_lora_manager = mocker.Mock() + wrapper.worker.load_model = mocker.Mock() custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} @@ -453,32 +459,30 @@ class TestCustomPipelineWorkerExtension: # Verify load_model was still called wrapper.worker.load_model.assert_called_once() - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_custom_pipeline_args_initialization(self, mock_worker_init, mock_od_config): + def test_custom_pipeline_args_initialization(self, mocker: MockerFixture, mock_od_config): """Test initialization with custom_pipeline_args calls re_init_pipeline.""" custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} - with patch.object(DiffusionWorker, "__init__", return_value=None): - with patch.object(WorkerWrapperBase, "_prepare_worker_class") as mock_prepare: - # Create a mock worker class with re_init_pipeline - mock_worker_class = Mock() - mock_worker_instance = Mock() - mock_worker_instance.re_init_pipeline = Mock() - mock_worker_class.return_value = mock_worker_instance - mock_prepare.return_value = mock_worker_class + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + mock_prepare = mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class") + # Create a mock worker class with re_init_pipeline + mock_worker_class = mocker.Mock() + mock_worker_instance = mocker.Mock() + mock_worker_instance.re_init_pipeline = mocker.Mock() + mock_worker_class.return_value = mock_worker_instance + mock_prepare.return_value = mock_worker_class - _ = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - custom_pipeline_args=custom_args, - ) + _ = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + custom_pipeline_args=custom_args, + ) - # Verify re_init_pipeline was called with custom_pipeline_args - mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args) + # Verify re_init_pipeline was called with custom_pipeline_args + mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args) - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_custom_pipeline_with_explicit_extension(self, mock_worker_init, mock_od_config): + def test_custom_pipeline_with_explicit_extension(self, mocker: MockerFixture, mock_od_config): """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided.""" class CustomExtension: @@ -490,24 +494,25 @@ class TestCustomPipelineWorkerExtension: custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} - with patch.object(DiffusionWorker, "__init__", return_value=None): - wrapper = WorkerWrapperBase( - gpu_id=0, - od_config=mock_od_config, - base_worker_class=DiffusionWorker, - worker_extension_cls=CustomExtension, - custom_pipeline_args=custom_args, - ) + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=CustomExtension, + custom_pipeline_args=custom_args, + ) - # Should still have the explicitly provided extension - assert CustomExtension in wrapper.worker.__class__.__bases__ - assert hasattr(wrapper.worker, "custom_extension_method") + # Should still have the explicitly provided extension + assert CustomExtension in wrapper.worker.__class__.__bases__ + assert hasattr(wrapper.worker, "custom_extension_method") - @patch("torch.cuda.empty_cache") - @patch("gc.collect") - @patch.object(DiffusionWorker, "__init__", return_value=None) - def test_re_init_pipeline_multiple_calls(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + def test_re_init_pipeline_multiple_calls(self, mocker: MockerFixture, mock_od_config): """Test calling re_init_pipeline multiple times.""" + mocker.patch("torch.cuda.empty_cache") + mocker.patch("gc.collect") + mocker.patch.object(DiffusionWorker, "__init__", return_value=None) + wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, @@ -516,13 +521,13 @@ class TestCustomPipelineWorkerExtension: ) # Setup mock model_runner - mock_model_runner = Mock() - mock_pipeline1 = Mock() - mock_pipeline2 = Mock() + mock_model_runner = mocker.Mock() + mock_pipeline1 = mocker.Mock() + mock_pipeline2 = mocker.Mock() mock_model_runner.pipeline = mock_pipeline1 wrapper.worker.model_runner = mock_model_runner - wrapper.worker.init_lora_manager = Mock() - wrapper.worker.load_model = Mock() + wrapper.worker.init_lora_manager = mocker.Mock() + wrapper.worker.load_model = mocker.Mock() # First call custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} diff --git a/tests/distributed/omni_connectors/test_adapter_and_flow.py b/tests/distributed/omni_connectors/test_adapter_and_flow.py index 2dc472393..c2be5003c 100644 --- a/tests/distributed/omni_connectors/test_adapter_and_flow.py +++ b/tests/distributed/omni_connectors/test_adapter_and_flow.py @@ -1,19 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from unittest.mock import MagicMock - import pytest +from pytest_mock import MockerFixture from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector, try_send_via_connector from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig from vllm_omni.distributed.omni_connectors.utils.initialization import get_connectors_config_for_stage +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @pytest.fixture -def mock_objects(): - return {"connector": MagicMock(), "metrics": MagicMock(), "queue_fn": MagicMock()} +def mock_objects(mocker: MockerFixture): + return {"connector": mocker.MagicMock(), "metrics": mocker.MagicMock(), "queue_fn": mocker.MagicMock()} def test_send_success(mock_objects): @@ -138,7 +139,7 @@ def test_recv_no_connector(): assert inputs is None -def test_shm_connector_flow(): +def test_shm_connector_flow(mocker: MockerFixture): """ Verify the full flow: Send -> Adapter -> Connector -> Adapter -> Recv. Using real SharedMemoryConnector (inline mode for simplicity). @@ -161,7 +162,7 @@ def test_shm_connector_flow(): def mock_submit(payload): queue_capture.append(payload) - mock_metrics = MagicMock() + mock_metrics = mocker.MagicMock() # 3. Send success = try_send_via_connector( @@ -217,7 +218,7 @@ def test_get_connectors_for_stage(): assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2" -def test_recv_with_missing_metadata(): +def test_recv_with_missing_metadata(mocker: MockerFixture): """Test recv when queue payload is malformed (missing metadata).""" # Connector expects metadata but task doesn't have it task = { @@ -226,7 +227,7 @@ def test_recv_with_missing_metadata(): "from_stage": "0", # Missing "connector_metadata" } - mock_conn = MagicMock() + mock_conn = mocker.MagicMock() # If get is called with None metadata, connector usually handles it or adapter handles exception mock_conn.get.side_effect = Exception("Get failed") diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index a7b4b2d01..7b416be72 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -1,14 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from unittest.mock import MagicMock import pytest +from pytest_mock import MockerFixture from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer +# pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_basic_serialization(): """Test basic msgpack serialization.""" @@ -79,19 +81,19 @@ def test_put_get_inline(shm_connector): assert size == ret_size -def test_put_get_shm(shm_connector, monkeypatch): +def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.MonkeyPatch): """Test SHM transfer logic for large data (Mocked).""" # Create data larger than 100 bytes data = {"large": "x" * 200} # Mock SHM return values mock_handle = {"name": "test_shm", "size": 200} - mock_write = MagicMock(return_value=mock_handle) + mock_write = mocker.MagicMock(return_value=mock_handle) monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write) # When reading, return the serialized bytes of the data serialized_data = shm_connector.serialize_obj(data) - mock_read = MagicMock(return_value=serialized_data) + mock_read = mocker.MagicMock(return_value=serialized_data) monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_read_bytes", mock_read) # Put diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index a5ebc2227..1d8341d4d 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -4,10 +4,10 @@ import threading from collections import deque from types import SimpleNamespace -from unittest.mock import MagicMock import pytest import torch +from pytest_mock import MockerFixture from vllm.v1.request import RequestStatus from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase @@ -39,9 +39,9 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None) @pytest.fixture -def build_adapter(monkeypatch): +def build_adapter(monkeypatch, mocker: MockerFixture): def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2): - connector = MagicMock() + connector = mocker.MagicMock() connector.stage_id = stage_id connector.get.return_value = None connector.put.return_value = (True, 1, {}) diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py index 84321f739..968109976 100755 --- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py @@ -7,7 +7,6 @@ These tests do NOT require Mooncake or RDMA environment. """ import threading -import unittest import pytest import torch @@ -22,7 +21,7 @@ pytestmark = [pytest.mark.cpu, pytest.mark.parallel] @pytest.mark.core_model -class TestBufferAllocator(unittest.TestCase): +class TestBufferAllocator: """Unit tests for BufferAllocator.""" def test_basic_alloc_free(self): @@ -30,15 +29,15 @@ class TestBufferAllocator(unittest.TestCase): allocator = BufferAllocator(total_size=4096, alignment=64) offset1 = allocator.alloc(512) - self.assertEqual(offset1, 0) + assert offset1 == 0 offset2 = allocator.alloc(512) - self.assertGreater(offset2, 0) + assert offset2 > 0 # Free first block, should be reusable allocator.free(offset1, 512) offset3 = allocator.alloc(512) - self.assertEqual(offset3, 0) + assert offset3 == 0 def test_alignment(self): """Verify allocation respects alignment.""" @@ -47,22 +46,22 @@ class TestBufferAllocator(unittest.TestCase): _offset1 = allocator.alloc(100) offset2 = allocator.alloc(100) - self.assertEqual(offset2 % 128, 0) - self.assertEqual(offset2, 128) + assert offset2 % 128 == 0 + assert offset2 == 128 def test_exhaustion_and_recovery(self): """Test that full allocation fails, then succeeds after free.""" allocator = BufferAllocator(total_size=1024, alignment=64) offset = allocator.alloc(1024) - self.assertEqual(offset, 0) + assert offset == 0 - with self.assertRaises(MemoryError): + with pytest.raises(MemoryError): allocator.alloc(64) allocator.free(offset, 1024) offset2 = allocator.alloc(1024) - self.assertEqual(offset2, 0) + assert offset2 == 0 def test_thread_safety(self): """Verify allocator is thread-safe under concurrent access.""" @@ -84,10 +83,10 @@ class TestBufferAllocator(unittest.TestCase): for t in threads: t.join() - self.assertEqual(len(errors), 0, f"Thread safety errors: {errors}") + assert len(errors) == 0, f"Thread safety errors: {errors}" -class TestAllocatorInvariants(unittest.TestCase): +class TestAllocatorInvariants: """ Defensive invariant tests for BufferAllocator: double-free, partial overlap corruption, and adjacent-block merging. @@ -106,7 +105,7 @@ class TestAllocatorInvariants(unittest.TestCase): allocator.free(offset, 256) # Pool should still be consistent: allocate full size back offset2 = allocator.alloc(4096) - self.assertEqual(offset2, 0) + assert offset2 == 0 @pytest.mark.slow def test_double_free_after_merge_is_safe(self): @@ -123,7 +122,7 @@ class TestAllocatorInvariants(unittest.TestCase): allocator.free(a, 64) # should not raise # Pool should still be fully usable offset = allocator.alloc(4096) - self.assertEqual(offset, 0) + assert offset == 0 @pytest.mark.slow def test_partial_overlap_raises_corruption(self): @@ -133,7 +132,7 @@ class TestAllocatorInvariants(unittest.TestCase): b = allocator.alloc(128) allocator.free(a, 128) # [0, 128) is now free # Try to free a region that starts inside [0,128) but extends beyond - with self.assertRaises(RuntimeError): + with pytest.raises(RuntimeError): allocator.free(64, 128) # [64, 192) overlaps with free [0, 128) # b is still allocated; freeing b should be fine @@ -154,7 +153,7 @@ class TestAllocatorInvariants(unittest.TestCase): # covering at least 192 bytes (3 * 64). # Verify by allocating a contiguous block of 192 bytes. offset = allocator.alloc(192) - self.assertEqual(offset, 0, "Adjacent blocks were not merged properly") + assert offset == 0, "Adjacent blocks were not merged properly" @pytest.mark.slow def test_fragmentation_and_defrag(self): @@ -174,7 +173,7 @@ class TestAllocatorInvariants(unittest.TestCase): allocator.free(d, 64) # free blocks: [64, 128) and [192, 256) # Pool has two 64-byte holes; contiguous 128 is unavailable - with self.assertRaises(MemoryError): + with pytest.raises(MemoryError): allocator.alloc(128) allocator.free(a, 64) @@ -182,14 +181,16 @@ class TestAllocatorInvariants(unittest.TestCase): # After freeing everything, full pool should be available offset = allocator.alloc(256) - self.assertEqual(offset, 0) + assert offset == 0 @pytest.mark.core_model -class TestManagedBuffer(unittest.TestCase): +class TestManagedBuffer: """Unit tests for ManagedBuffer.""" - def setUp(self): + @pytest.fixture(autouse=True) + def _setup(self): + # automatically invoked for every test method in the class self.allocator = BufferAllocator(total_size=4096, alignment=64) self.pool = torch.zeros(4096, dtype=torch.uint8) @@ -203,12 +204,12 @@ class TestManagedBuffer(unittest.TestCase): self.pool[offset : offset + 64] = src.view(torch.uint8) # Raw uint8 view - self.assertEqual(buf.tensor.shape[0], 64) + assert buf.tensor.shape[0] == 64 # Typed view typed = buf.as_tensor(dtype=torch.float32, shape=(4, 4)) - self.assertEqual(typed.shape, (4, 4)) - self.assertTrue(torch.equal(typed.flatten(), src)) + assert typed.shape == (4, 4) + assert torch.equal(typed.flatten(), src) buf.release() @@ -217,14 +218,10 @@ class TestManagedBuffer(unittest.TestCase): offset = self.allocator.alloc(128) with ManagedBuffer(self.allocator, offset, 128, self.pool) as buf: - self.assertFalse(buf._released) + assert not buf._released - self.assertTrue(buf._released) + assert buf._released # Space should be reusable new_offset = self.allocator.alloc(128) - self.assertEqual(new_offset, offset) - - -if __name__ == "__main__": - unittest.main() + assert new_offset == offset diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py index 698804f88..7b376da62 100755 --- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py @@ -13,7 +13,6 @@ import os import subprocess import threading import time -import unittest import pytest import torch @@ -195,43 +194,43 @@ def _md5(tensor: torch.Tensor) -> str: # --------------------------------------------------------------------------- -@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available") -class TestBasicConnector(unittest.TestCase): +@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available") +class TestBasicConnector: """Verify connector initialization, put, cleanup, and health check.""" def test_initialization(self): port = _free_port() with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c: - self.assertNotEqual(c.rpc_port, 0) - self.assertEqual(c.pool_size, 1024 * 1024) - self.assertTrue(c.pool.is_pinned()) + assert c.rpc_port != 0 + assert c.pool_size == 1024 * 1024 + assert c.pool.is_pinned() health = c.health() - self.assertEqual(health["status"], "healthy") + assert health["status"] == "healthy" def test_put_tensor_bytes_object(self): """Put tensor / bytes / dict and verify metadata.""" port = _free_port() with MooncakeTransferEngineConnector(_connector_config(port)) as c: ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100)) - self.assertTrue(ok) - self.assertTrue(meta["is_fast_path"]) + assert ok + assert meta["is_fast_path"] ok, sz, meta = c.put("s0", "s1", "b", b"hello" * 100) - self.assertTrue(ok) - self.assertTrue(meta["is_fast_path"]) + assert ok + assert meta["is_fast_path"] ok, sz, meta = c.put("s0", "s1", "d", {"k": [1, 2, 3]}) - self.assertTrue(ok) - self.assertFalse(meta["is_fast_path"]) + assert ok + assert not meta["is_fast_path"] def test_cleanup_releases_buffer(self): port = _free_port() with MooncakeTransferEngineConnector(_connector_config(port)) as c: c.put("s0", "s1", "r1", torch.randn(100)) key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1") - self.assertIn(key, c._local_buffers) + assert key in c._local_buffers c.cleanup("r1", from_stage="s0", to_stage="s1") - self.assertNotIn(key, c._local_buffers) + assert key not in c._local_buffers def test_pool_exhaustion_and_recovery(self): """Fill pool, verify failure, free, verify recovery.""" @@ -247,7 +246,7 @@ class TestBasicConnector(unittest.TestCase): for rid in ids: c.cleanup(rid, from_stage="s", to_stage="s") ok, _, _ = c.put("s", "s", "recovery", torch.randn(1000)) - self.assertTrue(ok, "Pool recovery failed after cleanup") + assert ok, "Pool recovery failed after cleanup" # --------------------------------------------------------------------------- @@ -255,8 +254,8 @@ class TestBasicConnector(unittest.TestCase): # --------------------------------------------------------------------------- -@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available") -class TestEndToEnd(unittest.TestCase): +@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available") +class TestEndToEnd: """E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types.""" def _pair(self, pool_size=16 * 1024 * 1024): @@ -269,14 +268,14 @@ class TestEndToEnd(unittest.TestCase): try: orig = torch.randn(1024, 1024, dtype=torch.float32) ok, sz, meta = p.put("s0", "s1", "t1", orig) - self.assertTrue(ok) + assert ok time.sleep(0.5) result = c.get("s0", "s1", "t1", meta) - self.assertIsNotNone(result) + assert result is not None buf, _ = result - self.assertIsInstance(buf, ManagedBuffer) + assert isinstance(buf, ManagedBuffer) recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape) - self.assertTrue(torch.equal(recon, orig)) + assert torch.equal(recon, orig) buf.release() finally: p.close() @@ -287,15 +286,15 @@ class TestEndToEnd(unittest.TestCase): try: orig = b"Hello RDMA! " * 1000 ok, _, meta = p.put("s0", "s1", "b1", orig) - self.assertTrue(ok) + assert ok time.sleep(0.5) result = c.get("s0", "s1", "b1", meta) - self.assertIsNotNone(result) + assert result is not None buf, _ = result recv = buf.to_bytes() if hasattr(buf, "to_bytes") else buf if hasattr(buf, "release"): buf.release() - self.assertEqual(recv, orig) + assert recv == orig finally: p.close() c.close() @@ -305,13 +304,13 @@ class TestEndToEnd(unittest.TestCase): try: orig = {"msg": "hello", "nums": [1, 2, 3], "nested": {"k": "v"}} ok, _, meta = p.put("s0", "s1", "o1", orig) - self.assertTrue(ok) - self.assertFalse(meta["is_fast_path"]) + assert ok + assert not meta["is_fast_path"] time.sleep(1.0) result = c.get("s0", "s1", "o1", meta) - self.assertIsNotNone(result) + assert result is not None recv, _ = result - self.assertEqual(recv, orig) + assert recv == orig finally: p.close() c.close() @@ -327,15 +326,15 @@ class TestEndToEnd(unittest.TestCase): sbuf.as_tensor(dtype=torch.float32, shape=shape).copy_(ref) ok, sz, meta = p.put("s0", "s1", "zc", sbuf) - self.assertTrue(ok) - self.assertTrue(meta.get("is_fast_path")) + assert ok + assert meta.get("is_fast_path") time.sleep(0.5) result = c.get("s0", "s1", "zc", meta) - self.assertIsNotNone(result) + assert result is not None rbuf, _ = result recon = rbuf.as_tensor(dtype=torch.float32, shape=shape) - self.assertTrue(torch.equal(recon, ref)) + assert torch.equal(recon, ref) rbuf.release() finally: p.close() @@ -348,13 +347,13 @@ class TestEndToEnd(unittest.TestCase): orig = torch.randn(5000, 5000, dtype=torch.float32) md5_send = _md5(orig) ok, _, meta = p.put("s0", "s1", "lg", orig) - self.assertTrue(ok) + assert ok time.sleep(1.0) result = c.get("s0", "s1", "lg", meta) - self.assertIsNotNone(result) + assert result is not None buf, _ = result recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape) - self.assertEqual(_md5(recon), md5_send) + assert _md5(recon) == md5_send buf.release() finally: p.close() @@ -374,23 +373,23 @@ class TestEndToEnd(unittest.TestCase): time.sleep(0.5) for name, data in cases: ok, _, meta = p.put("s0", "s1", name, data) - self.assertTrue(ok, f"put failed: {name}") + assert ok, f"put failed: {name}" time.sleep(0.3) result = c.get("s0", "s1", name, meta) - self.assertIsNotNone(result, f"get failed: {name}") + assert result is not None, f"get failed: {name}" recv, _ = result if isinstance(data, torch.Tensor): - self.assertIsInstance(recv, ManagedBuffer) + assert isinstance(recv, ManagedBuffer) recon = recv.as_tensor(dtype=data.dtype, shape=data.shape) - self.assertTrue(torch.equal(recon, data), f"mismatch: {name}") + assert torch.equal(recon, data), f"mismatch: {name}" recv.release() elif isinstance(data, bytes): got = recv.to_bytes() if hasattr(recv, "to_bytes") else recv if hasattr(recv, "release"): recv.release() - self.assertEqual(got, data) + assert got == data else: - self.assertEqual(recv, data) + assert recv == data finally: p.close() c.close() @@ -420,24 +419,24 @@ class TestEndToEnd(unittest.TestCase): t.start() for t in threads: t.join() - self.assertEqual(len(errors), 0, f"errors: {errors}") + assert len(errors) == 0, f"errors: {errors}" def test_auto_cleanup(self): """Producer buffer should be released after consumer get.""" p, c = self._pair() try: ok, _, meta = p.put("s0", "s1", "ac", torch.randn(100)) - self.assertTrue(ok) + assert ok key = MooncakeTransferEngineConnector._make_key("ac", "s0", "s1") - self.assertIn(key, p._local_buffers) + assert key in p._local_buffers time.sleep(0.5) result = c.get("s0", "s1", "ac", meta) - self.assertIsNotNone(result) + assert result is not None buf, _ = result if isinstance(buf, ManagedBuffer): buf.release() time.sleep(0.3) - self.assertNotIn(key, p._local_buffers) + assert key not in p._local_buffers finally: p.close() c.close() @@ -448,22 +447,22 @@ class TestEndToEnd(unittest.TestCase): # --------------------------------------------------------------------------- -@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available") -class TestLifecycle(unittest.TestCase): +@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available") +class TestLifecycle: """Close, context manager, double-close safety.""" def test_close_releases_resources(self): c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024)) c.put("s0", "s1", "x", torch.randn(100)) c.close() - self.assertTrue(c._stop_event.is_set()) - self.assertEqual(len(c._local_buffers), 0) + assert c._stop_event.is_set() + assert len(c._local_buffers) == 0 def test_context_manager(self): with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c: ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50)) - self.assertTrue(ok) - self.assertTrue(c._stop_event.is_set()) + assert ok + assert c._stop_event.is_set() def test_double_close_safe(self): c = MooncakeTransferEngineConnector(_connector_config(_free_port())) @@ -477,9 +476,9 @@ class TestLifecycle(unittest.TestCase): @pytest.mark.cuda -@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available") -@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available") -class TestGPUPool(unittest.TestCase): +@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestGPUPool: """GPU memory pool: initialization, put (CPU/GPU tensor), E2E transfer.""" @staticmethod @@ -488,18 +487,18 @@ class TestGPUPool(unittest.TestCase): def test_gpu_pool_init(self): with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c: - self.assertEqual(c.pool_device, "cuda:0") - self.assertTrue(c.pool.is_cuda) + assert c.pool_device == "cuda:0" + assert c.pool.is_cuda def test_gpu_pool_put_cpu_and_gpu_tensor(self): with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c: ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256)) - self.assertTrue(ok) - self.assertTrue(meta["is_fast_path"]) + assert ok + assert meta["is_fast_path"] ok, _, meta = c.put("s0", "s1", "d2d", torch.randn(256, 256, device="cuda:0")) - self.assertTrue(ok) - self.assertTrue(meta["is_fast_path"]) + assert ok + assert meta["is_fast_path"] def test_gpu_e2e_transfer(self): p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) @@ -507,14 +506,14 @@ class TestGPUPool(unittest.TestCase): try: orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0") ok, _, meta = p.put("s0", "s1", "ge", orig) - self.assertTrue(ok) + assert ok time.sleep(0.5) result = c.get("s0", "s1", "ge", meta) - self.assertIsNotNone(result) + assert result is not None buf, _ = result recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape) - self.assertTrue(recon.is_cuda) - self.assertTrue(torch.equal(recon.cpu(), orig.cpu())) + assert recon.is_cuda + assert torch.equal(recon.cpu(), orig.cpu()) buf.release() finally: p.close() @@ -527,8 +526,8 @@ class TestGPUPool(unittest.TestCase): @pytest.mark.slow -@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available") -class TestStressCorrectness(unittest.TestCase): +@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available") +class TestStressCorrectness: """ Slow but high-value regression tests: concurrent put+get with data integrity, large payload, edge cases, and sustained stress. @@ -581,20 +580,20 @@ class TestStressCorrectness(unittest.TestCase): for t in threads: t.join() - self.assertEqual(len(errors), 0, f"Producer errors: {errors}") - self.assertEqual(len(meta_map), num_workers) + assert len(errors) == 0, f"Producer errors: {errors}" + assert len(meta_map) == num_workers time.sleep(1.0) # Consumer: get each tensor and verify integrity for rid, meta in meta_map.items(): result = c.get("s0", "s1", rid, meta) - self.assertIsNotNone(result, f"get returned None for {rid}") + assert result is not None, f"get returned None for {rid}" buf, _ = result - self.assertIsInstance(buf, ManagedBuffer, f"{rid} not ManagedBuffer") + assert isinstance(buf, ManagedBuffer), f"{rid} not ManagedBuffer" recon = buf.as_tensor(dtype=torch.float32, shape=(512, 512)) recv_md5 = _md5(recon) - self.assertEqual(recv_md5, md5_sent[rid], f"MD5 mismatch for {rid}") + assert recv_md5 == md5_sent[rid], f"MD5 mismatch for {rid}" buf.release() finally: p.close() @@ -683,12 +682,8 @@ class TestStressCorrectness(unittest.TestCase): for t in threads: t.join() ct.join(timeout=60) - self.assertEqual(len(errors), 0, f"errors: {errors}") - self.assertEqual( - consumed_count[0], - num_items, - f"Consumer only processed {consumed_count[0]}/{num_items} items", - ) + assert len(errors) == 0, f"errors: {errors}" + assert consumed_count[0] == num_items, f"Consumer only processed {consumed_count[0]}/{num_items} items" finally: p.close() c.close() @@ -701,13 +696,13 @@ class TestStressCorrectness(unittest.TestCase): try: orig = torch.tensor([3.14], dtype=torch.float32) ok, _, meta = p.put("s0", "s1", "tiny", orig) - self.assertTrue(ok) + assert ok time.sleep(0.5) result = c.get("s0", "s1", "tiny", meta) - self.assertIsNotNone(result) + assert result is not None buf, _ = result recon = buf.as_tensor(dtype=torch.float32, shape=(1,)) - self.assertTrue(torch.equal(recon, orig)) + assert torch.equal(recon, orig) buf.release() finally: p.close() @@ -718,7 +713,7 @@ class TestStressCorrectness(unittest.TestCase): port = _free_port() with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c: ok, sz, meta = c.put("s0", "s1", "empty_b", b"") - self.assertFalse(ok, "Empty bytes should be rejected by connector") + assert not ok, "Empty bytes should be rejected by connector" # -- Large payload stress (500 MB) -- @@ -731,13 +726,13 @@ class TestStressCorrectness(unittest.TestCase): orig = torch.randn(side, side, dtype=torch.float32) md5_send = _md5(orig) ok, _, meta = p.put("s0", "s1", "500mb", orig) - self.assertTrue(ok) + assert ok time.sleep(2.0) result = c.get("s0", "s1", "500mb", meta) - self.assertIsNotNone(result) + assert result is not None buf, _ = result recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape) - self.assertEqual(_md5(recon), md5_send) + assert _md5(recon) == md5_send buf.release() finally: p.close() @@ -752,12 +747,8 @@ class TestStressCorrectness(unittest.TestCase): for i in range(50): rid = f"cycle_{i}" ok, _, _ = c.put("s0", "s1", rid, torch.randn(256, 256)) - self.assertTrue(ok, f"Put failed at iteration {i}") + assert ok, f"Put failed at iteration {i}" c.cleanup(rid, from_stage="s0", to_stage="s1") # After all cycles, pool should be fully recovered ok, _, _ = c.put("s0", "s1", "final", torch.randn(1024, 1024)) - self.assertTrue(ok, "Pool not fully recovered after rapid cycles") - - -if __name__ == "__main__": - unittest.main() + assert ok, "Pool not fully recovered after rapid cycles" diff --git a/tests/e2e/offline_inference/test_ovis_image.py b/tests/e2e/offline_inference/test_ovis_image.py index f1bc73817..41e21bca3 100644 --- a/tests/e2e/offline_inference/test_ovis_image.py +++ b/tests/e2e/offline_inference/test_ovis_image.py @@ -12,11 +12,11 @@ Strategy: and correctness of the model definition itself, independent of the pipeline mocks. """ -from unittest.mock import MagicMock, patch - import pytest import torch +from pytest_mock import MockerFixture +from tests.utils import hardware_test from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig # Mock the OvisImageTransformer2DModel to avoid complex init if needed, @@ -29,15 +29,15 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams @pytest.fixture -def mock_dependencies(monkeypatch): +def mock_dependencies(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): """ Mock external dependencies to avoid loading real models. """ device = get_local_device() # Mock Tokenizer - mock_tokenizer = MagicMock() - mock_tokenizer.return_value = MagicMock( + mock_tokenizer = mocker.MagicMock() + mock_tokenizer.return_value = mocker.MagicMock( input_ids=torch.zeros((1, 50), dtype=torch.long, device=device), attention_mask=torch.ones((1, 50), dtype=torch.long, device=device), ) @@ -45,13 +45,13 @@ def mock_dependencies(monkeypatch): mock_tokenizer.model_max_length = 1024 # Mock Text Encoder - mock_text_encoder = MagicMock() + mock_text_encoder = mocker.MagicMock() mock_text_encoder.dtype = torch.float32 # Output of text encoder must be on the same device as inputs (which are moved to execution_device) mock_text_encoder.return_value.last_hidden_state = torch.randn(1, 50, 32, device=device) # Mock VAE - mock_vae = MagicMock() + mock_vae = mocker.MagicMock() mock_vae.config.block_out_channels = [128, 256, 512, 512] # Scale factor 8 mock_vae.config.scale_factor_temporal = 1 mock_vae.config.scale_factor_spatial = 8 @@ -63,8 +63,8 @@ def mock_dependencies(monkeypatch): mock_vae.to.return_value = mock_vae # Mock Scheduler - mock_scheduler = MagicMock() - mock_scheduler.config = MagicMock() + mock_scheduler = mocker.MagicMock() + mock_scheduler.config = mocker.MagicMock() # Timesteps on device to match latents during denoising loop interaction if needed mock_scheduler.timesteps = torch.tensor([1.0, 0.5, 0.0], device=device) mock_scheduler.set_timesteps.return_value = None @@ -95,7 +95,7 @@ def mock_dependencies(monkeypatch): @pytest.fixture -def ovis_pipeline(mock_dependencies, monkeypatch): +def ovis_pipeline(mocker: MockerFixture, mock_dependencies, monkeypatch: pytest.MonkeyPatch): """ Creates an OvisImagePipeline instance with mocked components. """ @@ -122,8 +122,8 @@ def ovis_pipeline(mock_dependencies, monkeypatch): # Mock Transformer Layer separately to avoid full init # We patch OvisImageTransformer2DModel class in the module - mock_transformer_cls = MagicMock() - mock_transformer_instance = MagicMock() + mock_transformer_cls = mocker.MagicMock() + mock_transformer_instance = mocker.MagicMock() mock_transformer_instance.dtype = torch.float32 mock_transformer_instance.in_channels = 16 # Must be 16 so num_channel_latents=4, packed=16 # Forward return: noise prediction @@ -144,8 +144,8 @@ def ovis_pipeline(mock_dependencies, monkeypatch): # Initialize pipeline # We use a dummy model path check override - with patch("os.path.exists", return_value=True): - pipeline = OvisImagePipeline(od_config=od_config) + mocker.patch("os.path.exists", return_value=True) + pipeline = OvisImagePipeline(od_config=od_config) return pipeline @@ -223,9 +223,11 @@ def test_resolution_check(ovis_pipeline): assert output is not None -def test_real_transformer_init_and_forward(): +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +def test_real_transformer_init_and_forward(mocker: MockerFixture): """Test the real OvisImageTransformer2DModel initialization and forward pass for coverage.""" - from unittest.mock import patch from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel @@ -249,42 +251,42 @@ def test_real_transformer_init_and_forward(): # Mock distributed state for QKVParallelLinear initialization # We patch get_tp_group because get_tensor_model_parallel_rank calls it and asserts _TP is not None - mock_group = MagicMock() + mock_group = mocker.MagicMock() mock_group.rank_in_group = 0 mock_group.world_size = 1 - with patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group): - # Initialize real model - model = OvisImageTransformer2DModel( - od_config=od_config, - patch_size=1, - in_channels=16, - out_channels=16, - num_single_layers=1, - attention_head_dim=8, - num_attention_heads=2, - joint_attention_dim=32, - axes_dims_rope=(2, 2, 4), - ).to(device) + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group) + # Initialize real model + model = OvisImageTransformer2DModel( + od_config=od_config, + patch_size=1, + in_channels=16, + out_channels=16, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=32, + axes_dims_rope=(2, 2, 4), + ).to(device) - # Create dummy inputs - B, Seq, C = 1, 16, 16 - hidden_states = torch.randn(B, Seq, C, device=device) - encoder_hidden_states = torch.randn(B, 10, 32, device=device) # joint_attention_dim=32 - timestep = torch.tensor([1], device=device) - img_ids = torch.zeros(Seq, 3, device=device) - txt_ids = torch.zeros(10, 3, device=device) + # Create dummy inputs + B, Seq, C = 1, 16, 16 + hidden_states = torch.randn(B, Seq, C, device=device) + encoder_hidden_states = torch.randn(B, 10, 32, device=device) # joint_attention_dim=32 + timestep = torch.tensor([1], device=device) + img_ids = torch.zeros(Seq, 3, device=device) + txt_ids = torch.zeros(10, 3, device=device) - # Run forward - output = model( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - img_ids=img_ids, - txt_ids=txt_ids, - return_dict=False, - ) + # Run forward + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + return_dict=False, + ) - assert output is not None - assert isinstance(output, tuple) - assert output[0].shape == hidden_states.shape + assert output is not None + assert isinstance(output, tuple) + assert output[0].shape == hidden_states.shape diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 0c6479cce..36c58b3d8 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -10,11 +10,11 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py. import base64 import io from argparse import Namespace -from unittest.mock import AsyncMock, Mock import pytest from fastapi.testclient import TestClient from PIL import Image +from pytest_mock import MockerFixture from vllm import SamplingParams from vllm_omni.entrypoints.openai.image_api_utils import ( @@ -128,11 +128,11 @@ class FakeAsyncOmni: @pytest.fixture -def mock_async_diffusion(): +def mock_async_diffusion(mocker: MockerFixture): """Mock AsyncOmniDiffusion instance that returns fake images""" - mock = Mock() + mock = mocker.Mock() mock.is_running = True # For health endpoint - mock.check_health = AsyncMock() # For LLM mode health check + mock.check_health = mocker.AsyncMock() # For LLM mode health check async def generate(**kwargs): # Return n PIL images wrapped in result object @@ -143,7 +143,7 @@ def mock_async_diffusion(): images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)] return MockGenerationResult(images) - mock.generate = AsyncMock(side_effect=generate) + mock.generate = mocker.AsyncMock(side_effect=generate) return mock diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index a015b9b5f..b93592b39 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -6,27 +6,26 @@ Tests that standard OpenAI API parameters (max_tokens, temperature, etc.) are correctly applied to the comprehension stage while preserving YAML defaults. """ -from unittest.mock import MagicMock - import pytest +from pytest_mock import MockerFixture from vllm.sampling_params import SamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @pytest.fixture -def mock_comprehension_stage(): +def mock_comprehension_stage(mocker: MockerFixture): """Create a mock comprehension stage with is_comprehension=True.""" - stage = MagicMock() + stage = mocker.MagicMock() stage.is_comprehension = True stage.model_stage = "comprehension" return stage @pytest.fixture -def mock_other_stage(): +def mock_other_stage(mocker: MockerFixture): """Create a mock non-comprehension stage.""" - stage = MagicMock() + stage = mocker.MagicMock() stage.is_comprehension = False stage.model_stage = "other" return stage @@ -57,9 +56,15 @@ def default_other_params(): @pytest.fixture -def mock_engine_client(mock_comprehension_stage, mock_other_stage, default_comprehension_params, default_other_params): +def mock_engine_client( + mock_comprehension_stage, + mock_other_stage, + default_comprehension_params, + default_other_params, + mocker: MockerFixture, +): """Create mock engine client with stage_list and default_sampling_params_list.""" - engine_client = MagicMock() + engine_client = mocker.MagicMock() engine_client.stage_list = [mock_comprehension_stage, mock_other_stage] engine_client.default_sampling_params_list = [ default_comprehension_params, @@ -80,9 +85,9 @@ def serving_chat(mock_engine_client): @pytest.fixture -def mock_request(): +def mock_request(mocker: MockerFixture): """Create a mock request with all OpenAI sampling params set to None.""" - request = MagicMock() + request = mocker.MagicMock() # OpenAI standard sampling fields request.temperature = None request.top_p = None @@ -294,35 +299,35 @@ def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client): assert instance._get_comprehension_stage_index() == 0 -def test_get_comprehension_stage_index_finds_second_stage(): +def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture): """Test finding comprehension stage when it's at index 1.""" from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat instance = object.__new__(OmniOpenAIServingChat) - other = MagicMock() + other = mocker.MagicMock() other.is_comprehension = False - comprehension = MagicMock() + comprehension = mocker.MagicMock() comprehension.is_comprehension = True - instance.engine_client = MagicMock() + instance.engine_client = mocker.MagicMock() instance.engine_client.stage_list = [other, comprehension] assert instance._get_comprehension_stage_index() == 1 -def test_get_comprehension_stage_index_raises_when_not_found(): +def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixture): """Test that ValueError is raised when no comprehension stage exists.""" from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat instance = object.__new__(OmniOpenAIServingChat) - stage1 = MagicMock() + stage1 = mocker.MagicMock() stage1.is_comprehension = False - stage2 = MagicMock() + stage2 = mocker.MagicMock() stage2.is_comprehension = False - instance.engine_client = MagicMock() + instance.engine_client = mocker.MagicMock() instance.engine_client.stage_list = [stage1, stage2] with pytest.raises(ValueError, match="No comprehension stage"): diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 11d676839..37ee72972 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1,13 +1,13 @@ # tests/entrypoints/openai/test_serving_speech.py import logging from inspect import Signature, signature -from unittest.mock import MagicMock, patch import numpy as np import pytest import torch from fastapi import FastAPI from fastapi.testclient import TestClient +from pytest_mock import MockerFixture from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio, OpenAICreateSpeechRequest @@ -26,25 +26,24 @@ class TestAudioMixin: def audio_mixin(self): return AudioMixin() - def test_stereo_to_mono_conversion(self, audio_mixin): + def test_stereo_to_mono_conversion(self, audio_mixin, mocker: MockerFixture): stereo_tensor = np.random.rand(24000, 2).astype(np.float32) audio_obj = CreateAudio(audio_tensor=stereo_tensor) - with ( - patch.object( - audio_mixin, "_apply_speed_adjustment", side_effect=lambda tensor, speed, sr: (tensor, sr) - ) as mock_speed, - patch("soundfile.write") as _, - ): - audio_mixin.create_audio(audio_obj) + mock_speed = mocker.patch.object( + audio_mixin, "_apply_speed_adjustment", side_effect=lambda tensor, speed, sr: (tensor, sr) + ) + mocker.patch("soundfile.write") - # Check that the tensor passed to speed adjustment is mono - mock_speed.assert_called_once() - adjusted_tensor = mock_speed.call_args[0][0] - assert len(adjusted_tensor) == 24000 + audio_mixin.create_audio(audio_obj) - @patch("librosa.effects.time_stretch") - def test_speed_adjustment(self, mock_time_stretch, audio_mixin): + # Check that the tensor passed to speed adjustment is mono + mock_speed.assert_called_once() + adjusted_tensor = mock_speed.call_args[0][0] + assert len(adjusted_tensor) == 24000 + + def test_speed_adjustment(self, audio_mixin, mocker: MockerFixture): + mock_time_stretch = mocker.patch("librosa.effects.time_stretch") mock_time_stretch.return_value = np.zeros(12000) audio_tensor = np.random.rand(24000).astype(np.float32) @@ -53,8 +52,8 @@ class TestAudioMixin: mock_time_stretch.assert_called_with(y=audio_tensor, rate=2.0) assert adjusted_audio.shape == (12000,) - @patch("soundfile.write") - def test_unsupported_format_fallback(self, mock_write, audio_mixin, caplog): + def test_unsupported_format_fallback(self, audio_mixin, caplog, mocker: MockerFixture): + mock_write = mocker.patch("soundfile.write") audio_tensor = np.random.rand(24000).astype(np.float32) # Use a format that is not in the list of supported formats audio_obj = CreateAudio(audio_tensor=audio_tensor, response_format="vorbis") @@ -66,50 +65,50 @@ class TestAudioMixin: write_kwargs = mock_write.call_args.kwargs assert write_kwargs["format"] == "WAV" - def test_mono_audio_preservation(self, audio_mixin): + def test_mono_audio_preservation(self, audio_mixin, mocker: MockerFixture): """Test that mono (1D) audio tensors are processed correctly and passed to writer.""" mono_tensor = np.random.rand(24000).astype(np.float32) audio_obj = CreateAudio(audio_tensor=mono_tensor) - with patch("soundfile.write") as mock_write: - audio_mixin.create_audio(audio_obj) + mock_write = mocker.patch("soundfile.write") + audio_mixin.create_audio(audio_obj) - mock_write.assert_called_once() - # Verify the tensor passed to soundfile.write is the exact 1D tensor - output_tensor = mock_write.call_args[0][1] - assert output_tensor.ndim == 1 - assert output_tensor.shape == (24000,) - assert np.array_equal(output_tensor, mono_tensor) + mock_write.assert_called_once() + # Verify the tensor passed to soundfile.write is the exact 1D tensor + output_tensor = mock_write.call_args[0][1] + assert output_tensor.ndim == 1 + assert output_tensor.shape == (24000,) + assert np.array_equal(output_tensor, mono_tensor) - def test_stereo_audio_preservation(self, audio_mixin): + def test_stereo_audio_preservation(self, audio_mixin, mocker: MockerFixture): """Test that stereo (2D) audio tensors are processed correctly and preserved.""" stereo_tensor = np.random.rand(24000, 2).astype(np.float32) audio_obj = CreateAudio(audio_tensor=stereo_tensor) - with patch("soundfile.write") as mock_write: - audio_mixin.create_audio(audio_obj) + mock_write = mocker.patch("soundfile.write") + audio_mixin.create_audio(audio_obj) - mock_write.assert_called_once() - # Verify the tensor passed to soundfile.write is the exact 2D tensor - output_tensor = mock_write.call_args[0][1] - assert output_tensor.ndim == 2 - assert output_tensor.shape == (24000, 2) - assert np.array_equal(output_tensor, stereo_tensor) + mock_write.assert_called_once() + # Verify the tensor passed to soundfile.write is the exact 2D tensor + output_tensor = mock_write.call_args[0][1] + assert output_tensor.ndim == 2 + assert output_tensor.shape == (24000, 2) + assert np.array_equal(output_tensor, stereo_tensor) - def test_speed_adjustment_bypass(self, audio_mixin): + def test_speed_adjustment_bypass(self, audio_mixin, mocker: MockerFixture): """Test that speed=1.0 bypasses the expensive librosa time stretching.""" audio_tensor = np.random.rand(24000).astype(np.float32) - with patch("librosa.effects.time_stretch") as mock_time_stretch: - # speed=1.0 should return immediately without calling librosa - result, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=1.0, sample_rate=24000) + mock_time_stretch = mocker.patch("librosa.effects.time_stretch") + # speed=1.0 should return immediately without calling librosa + result, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=1.0, sample_rate=24000) - mock_time_stretch.assert_not_called() - assert np.array_equal(result, audio_tensor) + mock_time_stretch.assert_not_called() + assert np.array_equal(result, audio_tensor) - @patch("librosa.effects.time_stretch") - def test_speed_adjustment_stereo_handling(self, mock_time_stretch, audio_mixin): + def test_speed_adjustment_stereo_handling(self, audio_mixin, mocker: MockerFixture): """Test that speed adjustment is attempted on stereo inputs.""" + mock_time_stretch = mocker.patch("librosa.effects.time_stretch") stereo_tensor = np.random.rand(24000, 2).astype(np.float32) # Mock return value representing a sped-up version (half length) mock_time_stretch.return_value = np.zeros((12000, 2), dtype=np.float32) @@ -161,22 +160,22 @@ def create_mock_audio_output_for_test( @pytest.fixture -def test_app(): +def test_app(mocker: MockerFixture): # Mock the engine client - mock_engine_client = MagicMock() + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False async def mock_generate_fn(*args, **kwargs): yield create_mock_audio_output_for_test(request_id=kwargs.get("request_id")) - mock_engine_client.generate = MagicMock(side_effect=mock_generate_fn) + mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_fn) mock_engine_client.default_sampling_params_list = [{}] # Mock models to have an is_base_model method - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True - mock_request_logger = MagicMock() + mock_request_logger = mocker.MagicMock() speech_server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, @@ -186,7 +185,7 @@ def test_app(): # Patch the signature of create_speech to remove 'raw_request' for FastAPI route introspection original_create_speech = speech_server.create_speech - _ = MagicMock(side_effect=original_create_speech) + _ = mocker.MagicMock(side_effect=original_create_speech) sig = signature(original_create_speech) @@ -253,11 +252,13 @@ class TestSpeechAPI: response = client.post("/v1/audio/speech", json=payload) assert response.status_code == 422 # Unprocessable Entity - @patch("vllm_omni.entrypoints.openai.serving_speech.OmniOpenAIServingSpeech.create_audio") - def test_speed_parameter_is_used(self, mock_create_audio, test_app): + def test_speed_parameter_is_used(self, test_app, mocker: MockerFixture): + mock_create_audio = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.OmniOpenAIServingSpeech.create_audio" + ) client = TestClient(test_app) - mock_audio_response = MagicMock() + mock_audio_response = mocker.MagicMock() mock_audio_response.audio_data = b"dummy_audio" mock_audio_response.media_type = "audio/wav" mock_create_audio.return_value = mock_audio_response @@ -287,17 +288,17 @@ class TestTTSMethods: """Unit tests for TTS validation and parameter building.""" @pytest.fixture - def speech_server(self): - mock_engine_client = MagicMock() + def speech_server(self, mocker: MockerFixture): + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.stage_list = None mock_engine_client.tts_max_instructions_length = None - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True return OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, - request_logger=MagicMock(), + request_logger=mocker.MagicMock(), ) def test_is_tts_detection_no_stage(self, speech_server): @@ -306,25 +307,25 @@ class TestTTSMethods: assert speech_server._is_tts is False assert speech_server._tts_stage is None - def test_is_tts_detection_with_tts_stage(self): + def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture): """Test TTS model detection when TTS stage exists.""" - mock_engine_client = MagicMock() + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.tts_max_instructions_length = None # Create a TTS stage - mock_stage = MagicMock() + mock_stage = mocker.MagicMock() mock_stage.model_stage = "qwen3_tts" mock_stage.tts_args = {} mock_engine_client.stage_list = [mock_stage] - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, - request_logger=MagicMock(), + request_logger=mocker.MagicMock(), ) assert server._is_tts is True @@ -379,24 +380,24 @@ class TestTTSMethods: assert params["language"] == ["English"] assert params["task_type"] == ["CustomVoice"] - def test_load_supported_speakers(self): + def test_load_supported_speakers(self, mocker: MockerFixture): """Test _load_supported_speakers.""" - mock_engine_client = MagicMock() + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.stage_list = None # Mock talker_config with mixed-case speaker names - mock_talker_config = MagicMock() + mock_talker_config = mocker.MagicMock() mock_talker_config.spk_id = {"Ryan": 0, "Vivian": 1, "Aiden": 2} mock_engine_client.model_config.hf_config.talker_config = mock_talker_config - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, - request_logger=MagicMock(), + request_logger=mocker.MagicMock(), ) # Verify speakers are normalized to lowercase @@ -407,34 +408,34 @@ class TestTTSMethods: # Fixture creates server with no CLI override and no TTS stage assert speech_server._max_instructions_length == 500 - def test_max_instructions_length_cli_override(self): + def test_max_instructions_length_cli_override(self, mocker: MockerFixture): """Test CLI override (stored in engine_client) takes highest priority.""" - mock_engine_client = MagicMock() + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.stage_list = None # CLI override is stored in engine_client mock_engine_client.tts_max_instructions_length = 1000 - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, - request_logger=MagicMock(), + request_logger=mocker.MagicMock(), ) # Value is cached during __init__ assert server._max_instructions_length == 1000 - def test_max_instructions_length_stage_config(self): + def test_max_instructions_length_stage_config(self, mocker: MockerFixture): """Test stage config value is used when no CLI override.""" - mock_engine_client = MagicMock() + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.tts_max_instructions_length = None # No CLI override - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True # Mock stage with tts_args - mock_stage = MagicMock() + mock_stage = mocker.MagicMock() mock_stage.model_stage = "qwen3_tts" mock_stage.tts_args = {"max_instructions_length": 750} mock_engine_client.stage_list = [mock_stage] @@ -442,22 +443,22 @@ class TestTTSMethods: server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, - request_logger=MagicMock(), + request_logger=mocker.MagicMock(), ) # Value is cached during __init__ assert server._max_instructions_length == 750 - def test_max_instructions_length_cli_overrides_stage_config(self): + def test_max_instructions_length_cli_overrides_stage_config(self, mocker: MockerFixture): """Test CLI override (in engine_client) takes precedence over stage config.""" - mock_engine_client = MagicMock() + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False # CLI override stored in engine_client mock_engine_client.tts_max_instructions_length = 2000 - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True # Mock stage with tts_args - mock_stage = MagicMock() + mock_stage = mocker.MagicMock() mock_stage.model_stage = "qwen3_tts" mock_stage.tts_args = {"max_instructions_length": 750} mock_engine_client.stage_list = [mock_stage] @@ -465,25 +466,25 @@ class TestTTSMethods: server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, - request_logger=MagicMock(), + request_logger=mocker.MagicMock(), ) # CLI value (2000) should override stage config (750) assert server._max_instructions_length == 2000 - def test_validate_instructions_length_uses_cached_value(self): + def test_validate_instructions_length_uses_cached_value(self, mocker: MockerFixture): """Test instructions length validation uses cached _max_instructions_length.""" - mock_engine_client = MagicMock() + mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.stage_list = None # CLI override with max length of 10 characters mock_engine_client.tts_max_instructions_length = 10 - mock_models = MagicMock() + mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, - request_logger=MagicMock(), + request_logger=mocker.MagicMock(), ) # Verify cached value diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 6df304ce5..dc9414a67 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -5,18 +5,20 @@ Unit tests for OpenAI-compatible video generation endpoints. """ import io -from unittest.mock import patch import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from PIL import Image from pydantic import ValidationError +from pytest_mock import MockerFixture from vllm_omni.entrypoints.openai.api_server import router from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest, VideoResponseFormat from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class MockVideoResult: def __init__(self, videos): @@ -55,27 +57,27 @@ def _make_test_image_bytes(size=(64, 64)) -> bytes: return buf.getvalue() -def test_t2v_video_generation_form(test_client): +def test_t2v_video_generation_form(test_client, mocker: MockerFixture): fps_values = [] def _fake_encode(video, fps): fps_values.append(fps) return "Zg==" - with patch( + mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", side_effect=_fake_encode, - ): - response = test_client.post( - "/v1/videos", - data={ - "prompt": "A cat runs across the street.", - "size": "640x360", - "seconds": "2", - "fps": "12", - "n": "2", - }, - ) + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "A cat runs across the street.", + "size": "640x360", + "seconds": "2", + "fps": "12", + "n": "2", + }, + ) assert response.status_code == 200 data = response.json() @@ -92,18 +94,18 @@ def test_t2v_video_generation_form(test_client): assert fps_values == [12, 12] -def test_i2v_video_generation_form(test_client): +def test_i2v_video_generation_form(test_client, mocker: MockerFixture): image_bytes = _make_test_image_bytes((48, 32)) - with patch( + mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", return_value="Zg==", - ): - response = test_client.post( - "/v1/videos", - data={"prompt": "A bear playing with yarn."}, - files={"input_reference": ("input.png", image_bytes, "image/png")}, - ) + ) + response = test_client.post( + "/v1/videos", + data={"prompt": "A bear playing with yarn."}, + files={"input_reference": ("input.png", image_bytes, "image/png")}, + ) assert response.status_code == 200 data = response.json() @@ -119,24 +121,24 @@ def test_i2v_video_generation_form(test_client): assert input_image.size == (48, 32) -def test_seconds_defaults_fps_and_frames(test_client): +def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture): fps_values = [] def _fake_encode(video, fps): fps_values.append(fps) return "Zg==" - with patch( + mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", side_effect=_fake_encode, - ): - response = test_client.post( - "/v1/videos", - data={ - "prompt": "A bird flying.", - "seconds": "3", - }, - ) + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "A bird flying.", + "seconds": "3", + }, + ) assert response.status_code == 200 engine = test_client.app.state.openai_serving_video._engine_client @@ -146,18 +148,18 @@ def test_seconds_defaults_fps_and_frames(test_client): assert fps_values == [24] -def test_size_param_sets_width_height(test_client): - with patch( +def test_size_param_sets_width_height(test_client, mocker: MockerFixture): + mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", return_value="Zg==", - ): - response = test_client.post( - "/v1/videos", - data={ - "prompt": "size test", - "size": "320x240", - }, - ) + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "size test", + "size": "320x240", + }, + ) assert response.status_code == 200 engine = test_client.app.state.openai_serving_video._engine_client @@ -166,23 +168,23 @@ def test_size_param_sets_width_height(test_client): assert captured.height == 240 -def test_sampling_params_pass_through(test_client): - with patch( +def test_sampling_params_pass_through(test_client, mocker: MockerFixture): + mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", return_value="Zg==", - ): - response = test_client.post( - "/v1/videos", - data={ - "prompt": "param pass", - "num_inference_steps": "30", - "guidance_scale": "6.5", - "guidance_scale_2": "8.0", - "true_cfg_scale": "4.0", - "boundary_ratio": "0.7", - "flow_shift": "0.25", - }, - ) + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "param pass", + "num_inference_steps": "30", + "guidance_scale": "6.5", + "guidance_scale_2": "8.0", + "true_cfg_scale": "4.0", + "boundary_ratio": "0.7", + "flow_shift": "0.25", + }, + ) assert response.status_code == 200 engine = test_client.app.state.openai_serving_video._engine_client @@ -264,19 +266,19 @@ def test_invalid_n_raises_validation_error(test_client): ) -def test_negative_prompt_and_seed_pass_through(test_client): - with patch( +def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture): + mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", return_value="Zg==", - ): - response = test_client.post( - "/v1/videos", - data={ - "prompt": "snowy mountain", - "negative_prompt": "blurry", - "seed": "123", - }, - ) + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "snowy mountain", + "negative_prompt": "blurry", + "seed": "123", + }, + ) assert response.status_code == 200 engine = test_client.app.state.openai_serving_video._engine_client diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py index 8b7977d9c..29e6dcf6a 100644 --- a/tests/entrypoints/test_omni_diffusion.py +++ b/tests/entrypoints/test_omni_diffusion.py @@ -2,9 +2,9 @@ import uuid import warnings from queue import Empty, Queue from typing import Any -from unittest.mock import MagicMock import pytest +from pytest_mock import MockerFixture from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -152,7 +152,7 @@ class _FakeZmqQueue: class _FakeStage: """Lightweight Stage stub for multi-process pipeline version with queue support.""" - def __init__(self, config, stage_init_timeout: int = 300): + def __init__(self, mocker: MockerFixture, config, stage_init_timeout: int = 300): # Handle both dict and object configs if isinstance(config, dict): config = _FakeStageConfig(config) @@ -178,6 +178,10 @@ class _FakeStage: self._out_q = None self._proc = None # Mock process reference self._stage_init_timeout = max(0, int(stage_init_timeout)) + # NOTE: pytest mocker fixture object is injected into test functions and classes + # with specific naming (e.g. `class TestXXX`, `def test_XXX` ) automatically, + # but **must be provided** when initializing this helper cls + self._mocker = mocker def attach_queues(self, in_q, out_q): """Attach input and output queues.""" @@ -196,11 +200,11 @@ class _FakeStage: ): """Mock init_stage_worker: don't start real process, just send stage_ready message.""" # Create a mock process object - self._proc = MagicMock() - self._proc.start = MagicMock() - self._proc.join = MagicMock() - self._proc.is_alive = MagicMock(return_value=False) - self._proc.terminate = MagicMock() + self._proc = self._mocker.MagicMock() + self._proc.start = self._mocker.MagicMock() + self._proc.join = self._mocker.MagicMock() + self._proc.is_alive = self._mocker.MagicMock(return_value=False) + self._proc.terminate = self._mocker.MagicMock() # Send stage_ready message to output queue if self._out_q is not None: try: @@ -264,21 +268,23 @@ def fake_stage_config(): } -def _setup_engine_mocks(monkeypatch): - """Helper function to set up common engine mocks.""" - fake_engine = MagicMock() +def _setup_engine_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): + """Helper function to set up common engine mocks. + fixture objects `monkeypatch` and `mocker` must be passed in. + """ + fake_engine = mocker.MagicMock() # Add necessary attributes to fake_engine - fake_engine.tokenizer = MagicMock() + fake_engine.tokenizer = mocker.MagicMock() fake_engine.log_stats = False - fake_engine.vllm_config = MagicMock() - fake_engine.vllm_config.model_config = MagicMock() + fake_engine.vllm_config = mocker.MagicMock() + fake_engine.vllm_config.model_config = mocker.MagicMock() fake_engine.vllm_config.model_config.io_processor_plugin = None - fake_engine.get_supported_tasks = MagicMock(return_value=[]) - fake_engine.model_config = MagicMock() + fake_engine.get_supported_tasks = mocker.MagicMock(return_value=[]) + fake_engine.model_config = mocker.MagicMock() fake_engine.model_config.io_processor_plugin = None # Add registry with resolve_model_cls method - fake_registry = MagicMock() - fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch")) + fake_registry = mocker.MagicMock() + fake_registry.resolve_model_cls = mocker.MagicMock(return_value=(mocker.MagicMock(), "test_arch")) fake_engine.model_config.registry = fake_registry fake_engine.vllm_config.model_config.registry = fake_registry @@ -327,21 +333,23 @@ def _setup_engine_mocks(monkeypatch): ) -def _setup_multiprocessing_mocks(monkeypatch): - """Helper function to set up multiprocessing mocks.""" +def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): + """Helper function to set up multiprocessing mocks. + fixture objects `monkeypatch` and `mocker` must be passed in. + """ import multiprocessing as mp # Mock Process - fake_process_class = MagicMock() - fake_process_instance = MagicMock() - fake_process_instance.start = MagicMock() - fake_process_instance.join = MagicMock() - fake_process_instance.is_alive = MagicMock(return_value=False) - fake_process_instance.terminate = MagicMock() + fake_process_class = mocker.MagicMock() + fake_process_instance = mocker.MagicMock() + fake_process_instance.start = mocker.MagicMock() + fake_process_instance.join = mocker.MagicMock() + fake_process_instance.is_alive = mocker.MagicMock(return_value=False) + fake_process_instance.terminate = mocker.MagicMock() fake_process_class.return_value = fake_process_instance # Mock get_context to return a context with Queue that returns _FakeQueue - fake_ctx = MagicMock() + fake_ctx = mocker.MagicMock() fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize) fake_ctx.Process = fake_process_class @@ -365,7 +373,7 @@ def _setup_multiprocessing_mocks(monkeypatch): ) -def _setup_ipc_mocks(monkeypatch): +def _setup_ipc_mocks(monkeypatch: pytest.MonkeyPatch): """Helper function to set up IPC function mocks.""" # Mock _encode: simple serialization @@ -385,7 +393,7 @@ def _setup_ipc_mocks(monkeypatch): monkeypatch.setattr("vllm_omni.entrypoints.omni._set", _fake_set, raising=False) -def _setup_log_mocks(monkeypatch): +def _setup_log_mocks(monkeypatch: pytest.MonkeyPatch): """Helper function to set up logging and stats mocks.""" # Mock OrchestratorMetrics to be a simple class that doesn't require file operations @@ -416,7 +424,7 @@ def _setup_log_mocks(monkeypatch): ) -def _setup_connector_mocks(monkeypatch, omni_module=None): +def _setup_connector_mocks(monkeypatch, mocker, omni_module=None): """Helper function to set up connector mocks for stage-to-stage forwarding. If omni_module is provided, mocks directly on the module. Otherwise, uses string path. @@ -429,7 +437,7 @@ def _setup_connector_mocks(monkeypatch, omni_module=None): fake_connectors = {} # Add connectors for common edges (0->1, 1->2, etc.) for i in range(10): # Support up to 10 stages - fake_connectors[(str(i), str(i + 1))] = MagicMock() + fake_connectors[(str(i), str(i + 1))] = mocker.MagicMock() return None, fake_connectors if omni_module is not None: @@ -476,15 +484,15 @@ def _setup_connector_adapter_mock(monkeypatch, omni_module): @pytest.fixture(autouse=True) -def mock_get_config(monkeypatch): +def mock_get_config(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): """Auto-mock get_config and related model loading functions to avoid model path validation.""" # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports # This prevents ImportError when async_omni is imported (which happens via omni_stage) import sys - fake_tokenizer = MagicMock() - fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3]) - fake_tokenizer.decode = MagicMock(return_value="test") + fake_tokenizer = mocker.MagicMock() + fake_tokenizer.encode = mocker.MagicMock(return_value=[1, 2, 3]) + fake_tokenizer.decode = mocker.MagicMock(return_value="test") # Mock init_tokenizer_from_configs (used in async_omni) def _mock_init_tokenizer_from_configs(model_config=None, **kwargs): @@ -555,7 +563,7 @@ def mock_get_config(monkeypatch): setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) # Now mock get_config and other functions - fake_hf_config = MagicMock() + fake_hf_config = mocker.MagicMock() fake_hf_config.model_type = "qwen2_5_omni" def _mock_get_config(model, **kwargs): @@ -597,7 +605,9 @@ def mock_get_config(monkeypatch): ) -def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config): +def test_initialize_stage_configs_called_when_none( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config +): """Test that stage configs are auto-loaded when stage_configs_path is None.""" def _fake_loader( @@ -627,11 +637,11 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi del sys.modules[module_name] # Set up mocks - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) # Mock load_and_resolve_stage_configs monkeypatch.setattr( @@ -643,7 +653,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi # Replace OmniStage monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) @@ -651,7 +661,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi import vllm_omni.entrypoints.omni as omni_module # Patch the imported class in the module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) from vllm_omni.entrypoints.omni import Omni @@ -672,7 +682,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi assert len(omni._stages_ready) == 2 -def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): +def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that generate raises ValueError when sampling_params_list length doesn't match.""" def _fake_loader( @@ -693,11 +703,11 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -706,13 +716,13 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) from vllm_omni.entrypoints.omni import Omni @@ -722,7 +732,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): omni.generate(prompts=["hi"], sampling_params_list=[]) -def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): +def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test multi-stage generation pipeline with queue polling.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 @@ -750,8 +760,8 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) @@ -762,16 +772,16 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -828,7 +838,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): assert omni.stage_list[1].process_engine_inputs([], []) is not None -def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config): +def test_generate_pipeline_with_batch_input(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test single-stage generation pipeline with multiple inputs in one batch.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 @@ -856,8 +866,8 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) @@ -868,16 +878,16 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -946,7 +956,9 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config): assert len(outputs) == 2 -def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): +def test_generate_no_final_output_returns_empty( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config +): """Test that generate returns empty list when all stages have final_output=False.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 @@ -975,8 +987,8 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) @@ -987,16 +999,16 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -1037,7 +1049,9 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): assert outputs == [] -def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config): +def test_generate_sampling_params_none_use_default( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config +): """Test that generate uses default sampling params when sampling_params_list is None.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 @@ -1066,8 +1080,8 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) @@ -1078,16 +1092,16 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -1121,7 +1135,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi omni.generate(prompts=["p"], sampling_params_list=None) -def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): +def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that _wait_for_stages_ready handles timeout correctly.""" def _fake_loader( @@ -1142,11 +1156,11 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -1158,21 +1172,21 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): class _FakeStageNoReady(_FakeStage): def init_stage_worker(self, *args, **kwargs): # Don't send stage_ready message - self._proc = MagicMock() - self._proc.start = MagicMock() - self._proc.join = MagicMock() - self._proc.is_alive = MagicMock(return_value=False) - self._proc.terminate = MagicMock() + self._proc = self._mocker.MagicMock() + self._proc.start = self._mocker.MagicMock() + self._proc.join = self._mocker.MagicMock() + self._proc.is_alive = self._mocker.MagicMock(return_value=False) + self._proc.terminate = self._mocker.MagicMock() monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) from vllm_omni.entrypoints.omni import Omni @@ -1183,7 +1197,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): assert len(omni._stages_ready) == 0 -def test_generate_handles_error_messages(monkeypatch, fake_stage_config): +def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that generate handles error messages from stages correctly.""" def _fake_loader( @@ -1204,11 +1218,11 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -1217,13 +1231,13 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -1264,7 +1278,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config): assert len(outputs) == 1 -def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config): +def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that close() sends shutdown signal to all input queues.""" def _fake_loader( @@ -1285,11 +1299,11 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -1298,14 +1312,14 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) from vllm_omni.entrypoints.omni import Omni diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index 49971f1b4..4ec98be90 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/tests/entrypoints/test_omni_llm.py @@ -2,9 +2,9 @@ import uuid import warnings from queue import Empty, Queue from typing import Any -from unittest.mock import MagicMock import pytest +from pytest_mock import MockerFixture from vllm import SamplingParams from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK @@ -123,7 +123,7 @@ class _FakeZmqQueue: class _FakeStage: """Lightweight Stage stub for multi-process pipeline version with queue support.""" - def __init__(self, config, stage_init_timeout: int = 300): + def __init__(self, mocker: MockerFixture, config, stage_init_timeout: int = 300): # Handle both dict and object configs if isinstance(config, dict): config = _FakeStageConfig(config) @@ -149,6 +149,8 @@ class _FakeStage: self._out_q = None self._proc = None # Mock process reference self._stage_init_timeout = max(0, int(stage_init_timeout)) + # NOTE: mocker fixture object **must be provided** + self._mocker = mocker def attach_queues(self, in_q, out_q): """Attach input and output queues.""" @@ -167,11 +169,11 @@ class _FakeStage: ): """Mock init_stage_worker: don't start real process, just send stage_ready message.""" # Create a mock process object - self._proc = MagicMock() - self._proc.start = MagicMock() - self._proc.join = MagicMock() - self._proc.is_alive = MagicMock(return_value=False) - self._proc.terminate = MagicMock() + self._proc = self._mocker.MagicMock() + self._proc.start = self._mocker.MagicMock() + self._proc.join = self._mocker.MagicMock() + self._proc.is_alive = self._mocker.MagicMock(return_value=False) + self._proc.terminate = self._mocker.MagicMock() # Send stage_ready message to output queue if self._out_q is not None: try: @@ -235,21 +237,21 @@ def fake_stage_config(): } -def _setup_engine_mocks(monkeypatch): +def _setup_engine_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): """Helper function to set up common engine mocks.""" - fake_engine = MagicMock() + fake_engine = mocker.MagicMock() # Add necessary attributes to fake_engine - fake_engine.tokenizer = MagicMock() + fake_engine.tokenizer = mocker.MagicMock() fake_engine.log_stats = False - fake_engine.vllm_config = MagicMock() - fake_engine.vllm_config.model_config = MagicMock() + fake_engine.vllm_config = mocker.MagicMock() + fake_engine.vllm_config.model_config = mocker.MagicMock() fake_engine.vllm_config.model_config.io_processor_plugin = None - fake_engine.get_supported_tasks = MagicMock(return_value=[]) - fake_engine.model_config = MagicMock() + fake_engine.get_supported_tasks = mocker.MagicMock(return_value=[]) + fake_engine.model_config = mocker.MagicMock() fake_engine.model_config.io_processor_plugin = None # Add registry with resolve_model_cls method - fake_registry = MagicMock() - fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch")) + fake_registry = mocker.MagicMock() + fake_registry.resolve_model_cls = mocker.MagicMock(return_value=(mocker.MagicMock(), "test_arch")) fake_engine.model_config.registry = fake_registry fake_engine.vllm_config.model_config.registry = fake_registry @@ -298,21 +300,21 @@ def _setup_engine_mocks(monkeypatch): ) -def _setup_multiprocessing_mocks(monkeypatch): +def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): """Helper function to set up multiprocessing mocks.""" import multiprocessing as mp # Mock Process - fake_process_class = MagicMock() - fake_process_instance = MagicMock() - fake_process_instance.start = MagicMock() - fake_process_instance.join = MagicMock() - fake_process_instance.is_alive = MagicMock(return_value=False) - fake_process_instance.terminate = MagicMock() + fake_process_class = mocker.MagicMock() + fake_process_instance = mocker.MagicMock() + fake_process_instance.start = mocker.MagicMock() + fake_process_instance.join = mocker.MagicMock() + fake_process_instance.is_alive = mocker.MagicMock(return_value=False) + fake_process_instance.terminate = mocker.MagicMock() fake_process_class.return_value = fake_process_instance # Mock get_context to return a context with Queue that returns _FakeQueue - fake_ctx = MagicMock() + fake_ctx = mocker.MagicMock() fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize) fake_ctx.Process = fake_process_class @@ -336,7 +338,7 @@ def _setup_multiprocessing_mocks(monkeypatch): ) -def _setup_ipc_mocks(monkeypatch): +def _setup_ipc_mocks(monkeypatch: pytest.MonkeyPatch): """Helper function to set up IPC function mocks.""" # Mock _encode: simple serialization @@ -356,7 +358,7 @@ def _setup_ipc_mocks(monkeypatch): monkeypatch.setattr("vllm_omni.entrypoints.omni._set", _fake_set, raising=False) -def _setup_log_mocks(monkeypatch): +def _setup_log_mocks(monkeypatch: pytest.MonkeyPatch): """Helper function to set up logging and stats mocks.""" # Mock OrchestratorMetrics to be a simple class that doesn't require file operations @@ -387,7 +389,7 @@ def _setup_log_mocks(monkeypatch): ) -def _setup_connector_mocks(monkeypatch, omni_module=None): +def _setup_connector_mocks(monkeypatch, mocker, omni_module=None): """Helper function to set up connector mocks for stage-to-stage forwarding. If omni_module is provided, mocks directly on the module. Otherwise, uses string path. @@ -400,7 +402,7 @@ def _setup_connector_mocks(monkeypatch, omni_module=None): fake_connectors = {} # Add connectors for common edges (0->1, 1->2, etc.) for i in range(10): # Support up to 10 stages - fake_connectors[(str(i), str(i + 1))] = MagicMock() + fake_connectors[(str(i), str(i + 1))] = mocker.MagicMock() return None, fake_connectors if omni_module is not None: @@ -447,15 +449,15 @@ def _setup_connector_adapter_mock(monkeypatch, omni_module): @pytest.fixture(autouse=True) -def mock_get_config(monkeypatch): +def mock_get_config(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): """Auto-mock get_config and related model loading functions to avoid model path validation.""" # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports # This prevents ImportError when async_omni is imported (which happens via omni_stage) import sys - fake_tokenizer = MagicMock() - fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3]) - fake_tokenizer.decode = MagicMock(return_value="test") + fake_tokenizer = mocker.MagicMock() + fake_tokenizer.encode = mocker.MagicMock(return_value=[1, 2, 3]) + fake_tokenizer.decode = mocker.MagicMock(return_value="test") # Mock init_tokenizer_from_configs (used in async_omni) def _mock_init_tokenizer_from_configs(model_config=None, **kwargs): @@ -526,7 +528,7 @@ def mock_get_config(monkeypatch): setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) # Now mock get_config and other functions - fake_hf_config = MagicMock() + fake_hf_config = mocker.MagicMock() fake_hf_config.model_type = "qwen2_5_omni" def _mock_get_config(model, **kwargs): @@ -568,7 +570,9 @@ def mock_get_config(monkeypatch): ) -def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config): +def test_initialize_stage_configs_called_when_none( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config +): """Test that stage configs are auto-loaded when stage_configs_path is None.""" def _fake_loader( @@ -598,11 +602,11 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi del sys.modules[module_name] # Set up mocks - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) # Mock load_and_resolve_stage_configs monkeypatch.setattr( @@ -614,7 +618,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi # Replace OmniStage monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) @@ -623,9 +627,9 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi # Patch the imported function and class in the module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) from vllm_omni.entrypoints.omni import Omni @@ -645,7 +649,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi assert len(omni._stages_ready) == 2 -def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): +def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that generate raises ValueError when sampling_params_list length doesn't match.""" def _fake_loader( @@ -668,11 +672,11 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -681,14 +685,14 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) from vllm_omni.entrypoints.omni import Omni @@ -697,7 +701,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config): omni.generate(prompts=["hi"], sampling_params_list=[]) -def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): +def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test multi-stage generation pipeline with queue polling.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 @@ -723,11 +727,11 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -736,16 +740,16 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -802,7 +806,9 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config): assert omni.stage_list[1].process_engine_inputs([], []) is not None -def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): +def test_generate_no_final_output_returns_empty( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config +): """Test that generate returns empty list when all stages have final_output=False.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 @@ -829,11 +835,11 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -842,16 +848,16 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -892,7 +898,9 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config): assert outputs == [] -def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config): +def test_generate_sampling_params_none_use_default( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config +): """Test that generate uses default sampling params when sampling_params_list is None.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 @@ -919,11 +927,11 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -932,16 +940,16 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) # Apply connector and adapter mocks after importing omni module - _setup_connector_mocks(monkeypatch, omni_module) + _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation @@ -975,7 +983,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi omni.generate(prompts=["p"], sampling_params_list=None) -def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): +def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that _wait_for_stages_ready handles timeout correctly.""" def _fake_loader( @@ -998,11 +1006,11 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -1014,22 +1022,22 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): class _FakeStageNoReady(_FakeStage): def init_stage_worker(self, *args, **kwargs): # Don't send stage_ready message - self._proc = MagicMock() - self._proc.start = MagicMock() - self._proc.join = MagicMock() - self._proc.is_alive = MagicMock(return_value=False) - self._proc.terminate = MagicMock() + self._proc = self._mocker.MagicMock() + self._proc.start = self._mocker.MagicMock() + self._proc.join = self._mocker.MagicMock() + self._proc.is_alive = self._mocker.MagicMock(return_value=False) + self._proc.terminate = self._mocker.MagicMock() monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs)) from vllm_omni.entrypoints.omni import Omni @@ -1039,7 +1047,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config): assert len(omni._stages_ready) == 0 -def test_generate_handles_error_messages(monkeypatch, fake_stage_config): +def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that generate handles error messages from stages correctly.""" def _fake_loader( @@ -1062,11 +1070,11 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -1075,14 +1083,14 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) # Mock uuid.uuid4() to return a predictable value for request ID generation test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000") @@ -1122,7 +1130,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config): assert len(outputs) == 1 -def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config): +def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that close() sends shutdown signal to all input queues.""" def _fake_loader( @@ -1145,11 +1153,11 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config): if module_name in sys.modules: del sys.modules[module_name] - _setup_engine_mocks(monkeypatch) - _setup_multiprocessing_mocks(monkeypatch) + _setup_engine_mocks(monkeypatch, mocker) + _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) - _setup_connector_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", @@ -1158,14 +1166,14 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config): ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", - lambda cfg, **kwargs: _FakeStage(cfg, **kwargs), + lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) - monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs)) + monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) from vllm_omni.entrypoints.omni import Omni diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index ab7358d9f..5eb6154f0 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -1,8 +1,8 @@ import os import sys -from unittest.mock import MagicMock import pytest +from pytest_mock import MockerFixture from vllm_omni.entrypoints.stage_utils import set_stage_devices @@ -45,16 +45,18 @@ def _make_dummy_torch(call_log): return _Torch -def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"): - """Create a mock platform for testing.""" - mock_platform = MagicMock() +def _make_mock_platform(mocker, device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"): + """Create a mock platform for testing. + mocker object has to be passed in to utilize this helper function. + """ + mock_platform = mocker.MagicMock() mock_platform.device_type = device_type mock_platform.device_control_env_var = env_var return mock_platform @pytest.mark.usefixtures("clean_gpu_memory_between_tests") -def test_set_stage_devices_respects_logical_ids(monkeypatch): +def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): # Preserve an existing logical mapping and ensure devices "0,1" map through it. monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7") call_log: list[int] = [] @@ -62,7 +64,7 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch): monkeypatch.setitem(sys.modules, "torch", dummy_torch) # Mock the platform at the source module where it's defined - mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES") + mock_platform = _make_mock_platform(mocker, device_type="cuda", env_var="CUDA_VISIBLE_DEVICES") monkeypatch.setattr( "vllm_omni.platforms.current_omni_platform", mock_platform, @@ -74,7 +76,7 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch): @pytest.mark.usefixtures("clean_gpu_memory_between_tests") -def test_set_stage_devices_npu_platform(monkeypatch): +def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): """Test that set_stage_devices works correctly for NPU platform.""" monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5") call_log: list[int] = [] @@ -99,7 +101,7 @@ def test_set_stage_devices_npu_platform(monkeypatch): monkeypatch.setitem(sys.modules, "torch", _NpuTorch) # Mock NPU platform at the source module where it's defined - mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES") + mock_platform = _make_mock_platform(mocker, device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES") monkeypatch.setattr( "vllm_omni.platforms.current_omni_platform", mock_platform, diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index e18db40c1..11b26d8c3 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -2,9 +2,9 @@ from collections import Counter from dataclasses import dataclass -from unittest.mock import patch import pytest +from pytest_mock import MockerFixture from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs @@ -140,7 +140,7 @@ class TestFilterDictLikeObject: assert 42 in result assert (1, 2) in result - def test_dict_with_recursive_structure(self): + def test_dict_with_recursive_structure(self, mocker: MockerFixture): """Test filtering dict with recursive/complex nested structure.""" input_dict = { "level1": { @@ -152,15 +152,15 @@ class TestFilterDictLikeObject: "normal": "value", } - with patch("vllm_omni.entrypoints.utils.logger"): - result = _filter_dict_like_object(input_dict) + mocker.patch("vllm_omni.entrypoints.utils.logger") + result = _filter_dict_like_object(input_dict) - # Normal key should exist - assert "normal" in result - # Level1 should exist - assert "level1" in result + # Normal key should exist + assert "normal" in result + # Level1 should exist + assert "level1" in result - def test_integration_with_convert_dataclasses(self): + def test_integration_with_convert_dataclasses(self, mocker: MockerFixture): """Test that _filter_dict_like_object integrates properly with _convert_dataclasses_to_dict.""" @dataclass @@ -174,32 +174,32 @@ class TestFilterDictLikeObject: "normal": "value", } - with patch("vllm_omni.entrypoints.utils.logger"): - result = _filter_dict_like_object(input_dict) + mocker.patch("vllm_omni.entrypoints.utils.logger") + result = _filter_dict_like_object(input_dict) - # Callable should be filtered - assert "func" not in result - # Config should be converted to dict - assert "config" in result - assert "normal" in result + # Callable should be filtered + assert "func" not in result + # Config should be converted to dict + assert "config" in result + assert "normal" in result class TestConvertDataclassesToDict: """Test suite for _convert_dataclasses_to_dict function.""" - def test_uses_filter_dict_like_object(self): + def test_uses_filter_dict_like_object(self, mocker: MockerFixture): """Test that _convert_dataclasses_to_dict uses _filter_dict_like_object for dicts.""" input_dict = { "normal": "value", "callable": lambda x: x, } - with patch("vllm_omni.entrypoints.utils.logger"): - result = _convert_dataclasses_to_dict(input_dict) + mocker.patch("vllm_omni.entrypoints.utils.logger") + result = _convert_dataclasses_to_dict(input_dict) - # Callable should be filtered out by _filter_dict_like_object - assert "normal" in result - assert "callable" not in result + # Callable should be filtered out by _filter_dict_like_object + assert "normal" in result + assert "callable" not in result class TestFilterDataclassKwargs: diff --git a/tests/test_outputs.py b/tests/test_outputs.py index c0c6201f6..8da19a498 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for OmniRequestOutput class.""" -from unittest.mock import Mock - +import pytest from PIL import Image +from pytest_mock import MockerFixture from vllm_omni.outputs import OmniRequestOutput +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class TestOmniRequestOutput: """Tests for OmniRequestOutput class.""" @@ -28,17 +30,17 @@ class TestOmniRequestOutput: assert output.is_diffusion_output assert output.num_images == 1 - def test_from_pipeline(self): + def test_from_pipeline(self, mocker: MockerFixture): """Test creating output from pipeline stage.""" - mock_request_output = Mock() + mock_request_output = mocker.Mock() mock_request_output.request_id = "pipeline-123" mock_request_output.prompt_token_ids = [1, 2, 3] - mock_request_output.outputs = [Mock()] + mock_request_output.outputs = [mocker.Mock()] mock_request_output.encoder_prompt_token_ids = None mock_request_output.prompt_logprobs = None mock_request_output.num_cached_tokens = 10 mock_request_output.kv_transfer_params = None - mock_request_output.multimodal_output = {"image": Mock()} + mock_request_output.multimodal_output = {"image": mocker.Mock()} output = OmniRequestOutput.from_pipeline( stage_id=0, @@ -51,9 +53,9 @@ class TestOmniRequestOutput: assert output.final_output_type == "text" assert output.is_pipeline_output - def test_prompt_token_ids_property(self): + def test_prompt_token_ids_property(self, mocker: MockerFixture): """Test prompt_token_ids property for streaming compatibility.""" - mock_request_output = Mock() + mock_request_output = mocker.Mock() mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5] output = OmniRequestOutput.from_pipeline( @@ -73,10 +75,10 @@ class TestOmniRequestOutput: ) assert output.prompt_token_ids is None - def test_outputs_property(self): + def test_outputs_property(self, mocker: MockerFixture): """Test outputs property for chat completion compatibility.""" - mock_output = Mock() - mock_request_output = Mock() + mock_output = mocker.Mock() + mock_request_output = mocker.Mock() mock_request_output.outputs = [mock_output] output = OmniRequestOutput.from_pipeline( @@ -96,9 +98,9 @@ class TestOmniRequestOutput: ) assert output.outputs == [] - def test_encoder_prompt_token_ids_property(self): + def test_encoder_prompt_token_ids_property(self, mocker: MockerFixture): """Test encoder_prompt_token_ids property.""" - mock_request_output = Mock() + mock_request_output = mocker.Mock() mock_request_output.encoder_prompt_token_ids = [10, 20, 30] output = OmniRequestOutput.from_pipeline( @@ -109,9 +111,9 @@ class TestOmniRequestOutput: assert output.encoder_prompt_token_ids == [10, 20, 30] - def test_num_cached_tokens_property(self): + def test_num_cached_tokens_property(self, mocker: MockerFixture): """Test num_cached_tokens property.""" - mock_request_output = Mock() + mock_request_output = mocker.Mock() mock_request_output.num_cached_tokens = 42 output = OmniRequestOutput.from_pipeline( @@ -122,11 +124,12 @@ class TestOmniRequestOutput: assert output.num_cached_tokens == 42 - def test_multimodal_output_property(self): + def test_multimodal_output_property(self, mocker: MockerFixture): """Test multimodal_output property.""" - mock_request_output = Mock() - mock_audio = Mock() + mock_request_output = mocker.Mock() + mock_audio = mocker.Mock() expected_output = {"audio": mock_audio} + mock_request_output.outputs = [] mock_request_output.multimodal_output = expected_output output = OmniRequestOutput.from_pipeline( @@ -153,9 +156,9 @@ class TestOmniRequestOutput: assert result["num_images"] == 1 assert result["prompt"] == "a cat" - def test_to_dict_pipeline(self): + def test_to_dict_pipeline(self, mocker: MockerFixture): """Test to_dict for pipeline output.""" - mock_request_output = Mock() + mock_request_output = mocker.Mock() mock_request_output.request_id = "pipeline-123" output = OmniRequestOutput.from_pipeline( diff --git a/tests/worker/test_process_gpu_memory.py b/tests/worker/test_process_gpu_memory.py index d8e0b969d..c27da311d 100644 --- a/tests/worker/test_process_gpu_memory.py +++ b/tests/worker/test_process_gpu_memory.py @@ -1,52 +1,54 @@ """Tests for process-scoped GPU memory accounting.""" import os -from unittest import mock import pytest +from pytest_mock import MockerFixture + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class TestParseCudaVisibleDevices: - def test_empty(self): + def test_empty(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices - with mock.patch.dict(os.environ, {}, clear=True): - os.environ.pop("CUDA_VISIBLE_DEVICES", None) - assert parse_cuda_visible_devices() == [] + mocker.patch.dict(os.environ, {}, clear=True) + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + assert parse_cuda_visible_devices() == [] - with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): - assert parse_cuda_visible_devices() == [] + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}) + assert parse_cuda_visible_devices() == [] - def test_integer_indices(self): + def test_integer_indices(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices - with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"}): - assert parse_cuda_visible_devices() == [2, 3, 5] + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"}) + assert parse_cuda_visible_devices() == [2, 3, 5] - with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}): - assert parse_cuda_visible_devices() == [0] + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}) + assert parse_cuda_visible_devices() == [0] - def test_uuids(self): + def test_uuids(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices uuid1 = "GPU-12345678-1234-1234-1234-123456789abc" uuid2 = "GPU-87654321-4321-4321-4321-cba987654321" - with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"}): - assert parse_cuda_visible_devices() == [uuid1, uuid2] + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"}) + assert parse_cuda_visible_devices() == [uuid1, uuid2] - def test_mig_ids(self): + def test_mig_ids(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices mig1 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/0/0" mig2 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/1/0" - with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"}): - assert parse_cuda_visible_devices() == [mig1, mig2] + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"}) + assert parse_cuda_visible_devices() == [mig1, mig2] - def test_spaces(self): + def test_spaces(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices - with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "}): - assert parse_cuda_visible_devices() == [2, 3, 5] + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "}) + assert parse_cuda_visible_devices() == [2, 3, 5] class TestGetProcessGpuMemory: @@ -68,87 +70,79 @@ class TestGetProcessGpuMemory: del tensor torch.cuda.empty_cache() - def test_raises_on_invalid_device(self): + def test_raises_on_invalid_device(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory - with ( - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), - mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1), - ): - with pytest.raises(RuntimeError, match="Invalid GPU device"): - get_process_gpu_memory(5) + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit") + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown") + mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1) - def test_returns_zero_when_process_not_found(self): + with pytest.raises(RuntimeError, match="Invalid GPU device"): + get_process_gpu_memory(5) + + def test_returns_zero_when_process_not_found(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory - with ( - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), - mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex"), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]), - ): - memory = get_process_gpu_memory(0) - assert memory == 0 + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit") + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown") + mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8) + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex") + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]) - def test_uses_uuid_when_provided(self): + memory = get_process_gpu_memory(0) + assert memory == 0 + + def test_uses_uuid_when_provided(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory uuid = "GPU-12345678-1234-1234-1234-123456789abc" - mock_handle = mock.MagicMock() + mock_handle = mocker.MagicMock() - with ( - mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), - mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle) as mock_by_uuid, - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]), - ): - memory = get_process_gpu_memory(0) - assert memory == 0 - mock_by_uuid.assert_called_once_with(uuid) + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}) + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit") + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown") + mock_by_uuid = mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle) + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]) - def test_raises_on_invalid_uuid(self): + memory = get_process_gpu_memory(0) + assert memory == 0 + mock_by_uuid.assert_called_once_with(uuid) + + def test_raises_on_invalid_uuid(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory uuid = "GPU-invalid-uuid" - with ( - mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), - mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID")), - ): - with pytest.raises(RuntimeError, match="Failed to get NVML handle"): - get_process_gpu_memory(0) + mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}) + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit") + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown") + mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID")) - def test_returns_none_on_nvml_init_failure(self): + with pytest.raises(RuntimeError, match="Failed to get NVML handle"): + get_process_gpu_memory(0) + + def test_returns_none_on_nvml_init_failure(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory - with ( - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable")), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), - ): - result = get_process_gpu_memory(0) - assert result is None + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable")) + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown") + result = get_process_gpu_memory(0) + assert result is None class TestIsProcessScopedMemoryAvailable: - def test_returns_true_when_nvml_works(self): + def test_returns_true_when_nvml_works(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available - with ( - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), - mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), - ): - assert is_process_scoped_memory_available() is True + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit") + mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown") + assert is_process_scoped_memory_available() is True - def test_returns_false_when_nvml_fails(self): + def test_returns_false_when_nvml_fails(self, mocker: MockerFixture): from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available - with mock.patch( + mocker.patch( "vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable"), - ): - assert is_process_scoped_memory_available() is False + ) + assert is_process_scoped_memory_available() is False