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

[Misc] Migrate L1 tests to use pytest-mock (#1315)

Browse Source 
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed 2026-02-26 17:08:39 +08:00
1 parent 9913c00144
commit ed08a216c1
20 files changed
+977 -954
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_diffusion_worker.py
+78 -83
View File
@@ -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),
tests/diffusion/test_worker_wrapper_base.py
+170 -165
View File
@@ -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"}