wehub-resource-sync
320 lines
9.4 KiB
Python
320 lines
9.4 KiB
Python
import sys
|
|
import uuid
|
|
from typing import Any, Dict, List
|
|
from unittest.mock import MagicMock, call
|
|
|
|
import pytest
|
|
|
|
from ray.serve.api import deployment
|
|
from ray.serve.schema import (
|
|
CeleryAdapterConfig,
|
|
TaskProcessorAdapter,
|
|
TaskProcessorConfig,
|
|
TaskResult,
|
|
)
|
|
from ray.serve.task_consumer import task_consumer, task_handler
|
|
|
|
|
|
class MockTaskProcessorAdapter(TaskProcessorAdapter):
|
|
"""Mock adapter for testing task processor functionality."""
|
|
|
|
def __init__(self, config: TaskProcessorConfig):
|
|
self._config = config
|
|
self._start_consumer_received: bool = False
|
|
self._stop_consumer_received: bool = False
|
|
self.register_task_handle_mock = MagicMock()
|
|
|
|
def initialize(self, consumer_concurrency: int = 3):
|
|
pass
|
|
|
|
def register_task_handle(self, func, name=None):
|
|
self.register_task_handle_mock(func, name=name)
|
|
|
|
def enqueue_task_sync(
|
|
self, task_name, args=None, kwargs=None, **options
|
|
) -> TaskResult:
|
|
pass
|
|
|
|
def get_task_status_sync(self, task_id) -> TaskResult:
|
|
pass
|
|
|
|
def start_consumer(self, **kwargs):
|
|
self._start_consumer_received = True
|
|
|
|
def stop_consumer(self, timeout: float = 10.0):
|
|
self._stop_consumer_received = True
|
|
|
|
def cancel_task_sync(self, task_id) -> bool:
|
|
pass
|
|
|
|
def get_metrics_sync(self) -> Dict[str, Any]:
|
|
pass
|
|
|
|
def health_check_sync(self) -> List[Dict]:
|
|
pass
|
|
|
|
|
|
@pytest.fixture
|
|
def config():
|
|
"""Provides a mock TaskProcessorConfig."""
|
|
queue_name = f"test_queue_{uuid.uuid4().hex}"
|
|
return TaskProcessorConfig(
|
|
queue_name=queue_name,
|
|
adapter_config=CeleryAdapterConfig(
|
|
broker_url="fake://",
|
|
backend_url="fake://",
|
|
),
|
|
adapter=MockTaskProcessorAdapter,
|
|
)
|
|
|
|
|
|
class TestTaskHandlerDecorator:
|
|
"""Test the task_handler decorator."""
|
|
|
|
def _create_and_test_handler(self, decorator_args=None, expected_name=None):
|
|
"""Helper to create and test a task handler."""
|
|
mock = MagicMock()
|
|
|
|
if decorator_args is None:
|
|
|
|
@task_handler
|
|
def test_handler():
|
|
mock()
|
|
|
|
else:
|
|
|
|
@task_handler(**decorator_args)
|
|
def test_handler():
|
|
mock()
|
|
|
|
test_handler()
|
|
|
|
assert mock.call_count == 1
|
|
assert test_handler._task_name == expected_name
|
|
|
|
def test_task_handler_decorator_with_name(self):
|
|
self._create_and_test_handler(
|
|
decorator_args={"name": "my_task"}, expected_name="my_task"
|
|
)
|
|
|
|
def test_task_handler_decorator_without_name(self):
|
|
self._create_and_test_handler(expected_name="test_handler")
|
|
|
|
@pytest.mark.parametrize("invalid_name", ["", " ", 123])
|
|
def test_task_handler_decorator_invalid_name(self, invalid_name):
|
|
"""Test various invalid task names."""
|
|
with pytest.raises(
|
|
ValueError,
|
|
match=f"Task name must be a non-empty string, got {invalid_name}",
|
|
):
|
|
|
|
@task_handler(name=invalid_name)
|
|
def my_task_handler():
|
|
pass
|
|
|
|
def test_task_handler_on_callable_object_without_name_attr(self):
|
|
"""Test that AttributeError is raised for callables with no __name__."""
|
|
|
|
class MyCallable:
|
|
"""A simple callable class without a __name__ attribute on instances."""
|
|
|
|
def __call__(self):
|
|
pass
|
|
|
|
with pytest.raises(AttributeError):
|
|
task_handler(MyCallable())
|
|
|
|
|
|
class TestTaskConsumerDecorator:
|
|
"""Test the task_consumer decorator."""
|
|
|
|
def _verify_and_cleanup(self, instance, expected_calls=None):
|
|
"""Verify consumer and cleanup instance."""
|
|
instance.initialize_callable(5)
|
|
adapter = instance._adapter
|
|
assert adapter._start_consumer_received
|
|
|
|
if expected_calls is not None:
|
|
if expected_calls:
|
|
calls = [call(method, name=name) for method, name in expected_calls]
|
|
adapter.register_task_handle_mock.assert_has_calls(
|
|
calls, any_order=False
|
|
)
|
|
assert adapter.register_task_handle_mock.call_count == len(
|
|
expected_calls
|
|
)
|
|
else:
|
|
adapter.register_task_handle_mock.assert_not_called()
|
|
|
|
del instance
|
|
|
|
def _run_consumer_test(
|
|
self, config, consumer_class_factory, expected_calls_factory=None
|
|
):
|
|
"""Run a consumer test with factory functions."""
|
|
consumer_class = consumer_class_factory(config)
|
|
instance = consumer_class()
|
|
|
|
expected_calls = (
|
|
expected_calls_factory(instance) if expected_calls_factory else None
|
|
)
|
|
|
|
self._verify_and_cleanup(instance, expected_calls)
|
|
|
|
def test_task_consumer_basic(self, config):
|
|
"""Test basic functionality of the task_consumer decorator."""
|
|
|
|
def make_consumer(cfg):
|
|
@task_consumer(task_processor_config=cfg)
|
|
class MyConsumer:
|
|
@task_handler
|
|
def my_task(self):
|
|
pass
|
|
|
|
return MyConsumer
|
|
|
|
self._run_consumer_test(
|
|
config, make_consumer, lambda inst: [(inst.my_task, "my_task")]
|
|
)
|
|
|
|
def test_task_consumer_multiple_handlers(self, config):
|
|
"""Test with multiple task handlers."""
|
|
|
|
def make_consumer(cfg):
|
|
@task_consumer(task_processor_config=cfg)
|
|
class MyConsumer:
|
|
@task_handler
|
|
def task1(self):
|
|
pass
|
|
|
|
@task_handler
|
|
def task2(self):
|
|
pass
|
|
|
|
return MyConsumer
|
|
|
|
self._run_consumer_test(
|
|
config,
|
|
make_consumer,
|
|
lambda inst: [(inst.task1, "task1"), (inst.task2, "task2")],
|
|
)
|
|
|
|
def test_task_consumer_custom_names(self, config):
|
|
"""Test task handlers with and without custom names."""
|
|
|
|
def make_consumer(cfg):
|
|
@task_consumer(task_processor_config=cfg)
|
|
class MyConsumer:
|
|
@task_handler(name="custom_task")
|
|
def task1(self):
|
|
pass
|
|
|
|
@task_handler
|
|
def task2(self):
|
|
pass
|
|
|
|
return MyConsumer
|
|
|
|
self._run_consumer_test(
|
|
config,
|
|
make_consumer,
|
|
lambda inst: [(inst.task1, "custom_task"), (inst.task2, "task2")],
|
|
)
|
|
|
|
def test_task_consumer_init_args(self, config):
|
|
"""Test that __init__ arguments are passed correctly."""
|
|
|
|
@task_consumer(task_processor_config=config)
|
|
class MyConsumer:
|
|
def __init__(self, value):
|
|
self.value = value
|
|
|
|
instance = MyConsumer(value=42)
|
|
assert instance.value == 42
|
|
self._verify_and_cleanup(instance)
|
|
|
|
def test_task_consumer_no_handlers(self, config):
|
|
"""Test with a class that has no task handlers."""
|
|
|
|
def make_consumer(cfg):
|
|
@task_consumer(task_processor_config=cfg)
|
|
class MyConsumer:
|
|
def some_method(self):
|
|
pass
|
|
|
|
return MyConsumer
|
|
|
|
self._run_consumer_test(config, make_consumer, lambda inst: [])
|
|
|
|
def test_task_consumer_inheritance(self, config):
|
|
"""Test that inherited task handlers are registered."""
|
|
|
|
def make_consumer(cfg):
|
|
class BaseConsumer:
|
|
@task_handler
|
|
def base_task(self):
|
|
pass
|
|
|
|
@task_consumer(task_processor_config=cfg)
|
|
class DerivedConsumer(BaseConsumer):
|
|
@task_handler
|
|
def derived_task(self):
|
|
pass
|
|
|
|
return DerivedConsumer
|
|
|
|
self._run_consumer_test(
|
|
config,
|
|
make_consumer,
|
|
lambda inst: [
|
|
(inst.base_task, "base_task"),
|
|
(inst.derived_task, "derived_task"),
|
|
],
|
|
)
|
|
|
|
def test_task_consumer_no_args_decorator(self):
|
|
"""Test using @task_consumer without arguments raises TypeError."""
|
|
with pytest.raises(TypeError):
|
|
|
|
@task_consumer
|
|
class MyConsumer:
|
|
pass
|
|
|
|
|
|
def test_default_deployment_name_stays_same_with_task_consumer(config):
|
|
"""Test that the default deployment name is the class name when using task_consumer with serve.deployment."""
|
|
|
|
@deployment
|
|
@task_consumer(task_processor_config=config)
|
|
class MyTaskConsumer:
|
|
@task_handler
|
|
def my_task(self):
|
|
pass
|
|
|
|
# The deployment name should default to the class name
|
|
assert MyTaskConsumer.name == "MyTaskConsumer"
|
|
|
|
|
|
def test_task_consumer_preserves_metadata(config):
|
|
class OriginalConsumer:
|
|
"""Docstring for a task consumer."""
|
|
|
|
value: int
|
|
|
|
wrapped_cls = task_consumer(task_processor_config=config)(OriginalConsumer)
|
|
|
|
assert wrapped_cls.__name__ == OriginalConsumer.__name__
|
|
assert wrapped_cls.__qualname__ == OriginalConsumer.__qualname__
|
|
assert wrapped_cls.__module__ == OriginalConsumer.__module__
|
|
assert wrapped_cls.__doc__ == OriginalConsumer.__doc__
|
|
assert (
|
|
wrapped_cls.__annotations__["value"]
|
|
== OriginalConsumer.__annotations__["value"]
|
|
)
|
|
assert wrapped_cls.__annotations__["_adapter"] is TaskProcessorAdapter
|
|
assert getattr(wrapped_cls, "__wrapped__", None) is OriginalConsumer
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|