diff --git a/miles/utils/workers/rpc/common/metadata.py b/miles/utils/workers/rpc/common/metadata.py index a1642b50243..ef0fdf9ce5c 100644 --- a/miles/utils/workers/rpc/common/metadata.py +++ b/miles/utils/workers/rpc/common/metadata.py @@ -44,7 +44,10 @@ def collect_rpc_method_specs(worker_cls: type) -> dict[str, RpcMethodSpec]: continue if not callable(static_attr): continue - specs[name] = _build_method_spec(worker_cls=worker_cls, name=name, fn=inspect.unwrap(static_attr)) + specs[name] = _build_method_spec(worker_cls=worker_cls, name=name, attr=static_attr) + + if len(specs) == 0: + raise TypeError(f"{worker_cls.__name__} exposes no public rpc methods") return specs @@ -54,9 +57,26 @@ class _RpcConfig: concurrency_group: str -def _build_method_spec(*, worker_cls: type, name: str, fn: Callable[..., Any]) -> RpcMethodSpec: - config: _RpcConfig = getattr(fn, _RPC_CONFIG_ATTR, _RpcConfig(concurrency_group=DEFAULT_CONCURRENCY_GROUP)) - is_async = inspect.iscoroutinefunction(inspect.unwrap(fn)) +def _find_rpc_config(attr: Callable[..., Any]) -> _RpcConfig: + layer: Any = attr + while layer is not None: + config = getattr(layer, _RPC_CONFIG_ATTR, None) + if config is not None: + return config + layer = getattr(layer, "__wrapped__", None) + return _RpcConfig(concurrency_group=DEFAULT_CONCURRENCY_GROUP) + + +def _build_method_spec(*, worker_cls: type, name: str, attr: Callable[..., Any]) -> RpcMethodSpec: + fn = inspect.unwrap(attr) + if not inspect.isroutine(fn): + raise TypeError( + f"{worker_cls.__name__}.{name} is a public callable attribute but not a method, " + f"so it cannot be exposed over rpc; make it private or move it off the worker class" + ) + + config = _find_rpc_config(attr) + is_async = inspect.iscoroutinefunction(fn) if is_async and config.concurrency_group != DEFAULT_CONCURRENCY_GROUP: raise TypeError( f"{worker_cls.__name__}.{name} is async; concurrency groups only serialize sync methods, " @@ -66,11 +86,39 @@ def _build_method_spec(*, worker_cls: type, name: str, fn: Callable[..., Any]) - signature = inspect.signature(fn) hints = typing.get_type_hints(fn, include_extras=True) + parameters = list(signature.parameters.values()) + if len(parameters) == 0: + raise TypeError(f"{worker_cls.__name__}.{name} must take a receiver parameter for rpc exposure") + if parameters[0].name != "self": + raise TypeError( + f"{worker_cls.__name__}.{name} must name its receiver parameter 'self' for rpc exposure, " + f"got {parameters[0].name!r}; otherwise it would be silently dropped from the wire" + ) + if parameters[0].kind not in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD): + raise TypeError( + f"{worker_cls.__name__}.{name} must take its receiver parameter positionally for rpc exposure, " + f"got a {parameters[0].kind.description} parameter" + ) + query_fields: dict[str, Any] = {} - for param in list(signature.parameters.values())[1:]: + for param in parameters[1:]: + if param.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + inspect.Parameter.POSITIONAL_ONLY, + ): + raise TypeError( + f"{worker_cls.__name__}.{name} must not use *args/**kwargs or positional-only parameters " + f"for rpc exposure" + ) + if param.annotation is inspect.Parameter.empty: + raise TypeError(f"{worker_cls.__name__}.{name} parameter '{param.name}' must be type-annotated") default = ... if param.default is inspect.Parameter.empty else param.default query_fields[param.name] = (hints[param.name], default) + if signature.return_annotation is inspect.Signature.empty: + raise TypeError(f"{worker_cls.__name__}.{name} must have a return type annotation") + return RpcMethodSpec( name=name, concurrency_group=config.concurrency_group, diff --git a/miles/utils/workers/rpc/common/serialization.py b/miles/utils/workers/rpc/common/serialization.py index a26c59fe119..63aef8ba391 100644 --- a/miles/utils/workers/rpc/common/serialization.py +++ b/miles/utils/workers/rpc/common/serialization.py @@ -3,8 +3,11 @@ import dataclasses from typing import Any +import typing_extensions from pydantic import BaseModel, ConfigDict, TypeAdapter, create_model +_WIRE_CONFIG = ConfigDict(ser_json_bytes="base64", val_json_bytes="base64") + @dataclasses.dataclass(frozen=True) class RpcSerializer: @@ -13,8 +16,11 @@ class RpcSerializer: @classmethod def create(cls, *, query_model_name: str, query_fields: dict[str, Any], result_annotation: Any) -> RpcSerializer: - query_model = create_model(query_model_name, __config__=ConfigDict(extra="forbid"), **query_fields) - return cls(query_model=query_model, result_adapter=TypeAdapter(result_annotation)) + query_model = create_model( + query_model_name, __config__=ConfigDict(extra="forbid", **_WIRE_CONFIG), **query_fields + ) + result_config = None if _carries_own_config(result_annotation) else _WIRE_CONFIG + return cls(query_model=query_model, result_adapter=TypeAdapter(result_annotation, config=result_config)) def encode_query(self, kwargs: dict[str, Any]) -> dict[str, Any]: return self.query_model(**kwargs).model_dump(mode="json") @@ -27,3 +33,13 @@ def encode_result(self, result: Any) -> Any: def decode_result(self, payload: Any) -> Any: return self.result_adapter.validate_python(payload) + + +def _carries_own_config(annotation: Any) -> bool: + if not isinstance(annotation, type): + return False + return ( + issubclass(annotation, BaseModel) + or dataclasses.is_dataclass(annotation) + or typing_extensions.is_typeddict(annotation) + ) diff --git a/tests/fast/utils/workers/e2e/test_client_local_validation.py b/tests/fast/utils/workers/e2e/test_client_local_validation.py new file mode 100644 index 00000000000..cd4dfac5178 --- /dev/null +++ b/tests/fast/utils/workers/e2e/test_client_local_validation.py @@ -0,0 +1,97 @@ +import uuid + +import httpx +import pytest +from pydantic import ValidationError +from tests.fast.utils.workers.e2e.e2e_worker import E2eWorker + +from miles.utils.workers.rpc.client.handle import RpcWorkerHandle + + +class TestNoRequestIsSent: + async def test_unknown_method_is_an_attribute_error(self, dead_proxy, make_handle): + """A method the worker does not define fails before any request goes out.""" + handle = make_handle(dead_proxy) + with pytest.raises(AttributeError, match="no rpc method"): + _ = handle.no_such_method + assert dead_proxy.requests == [] + + async def test_missing_required_argument(self, dead_proxy, make_handle): + """A missing argument is caught locally, not by the server.""" + handle = make_handle(dead_proxy) + with pytest.raises(ValidationError): + await handle.demo_sync(a=1) + assert dead_proxy.requests == [] + + async def test_unknown_argument(self, dead_proxy, make_handle): + """An argument the method does not declare is caught locally.""" + handle = make_handle(dead_proxy) + with pytest.raises(ValidationError): + await handle.demo_sync(a=1, b=2, c=3) + assert dead_proxy.requests == [] + + async def test_wrong_argument_type(self, dead_proxy, make_handle): + """An uncoercible argument is caught locally.""" + handle = make_handle(dead_proxy) + with pytest.raises(ValidationError): + await handle.demo_sync(a="not-a-number", b=2) + assert dead_proxy.requests == [] + + async def test_positional_arguments_are_rejected(self, dead_proxy, make_handle): + """Calls are keyword-only, so positional arguments fail locally.""" + handle = make_handle(dead_proxy) + with pytest.raises(TypeError): + await handle.demo_sync(1, 2) + assert dead_proxy.requests == [] + + +class TestHandleConstruction: + def test_reserved_method_name_is_rejected(self): + """A worker whose method shadows a handle attribute is refused at construction.""" + + class Shadowing: + async def wait_ready(self, timeout: float) -> None: + pass + + with pytest.raises(TypeError, match="shadow"): + RpcWorkerHandle(Shadowing, server_url="http://127.0.0.1:9") + + def test_worker_without_public_methods_is_rejected(self): + """A worker with nothing to expose is refused.""" + + class Empty: + def _demo_hidden(self) -> int: + return 1 + + with pytest.raises(TypeError): + RpcWorkerHandle(Empty, server_url="http://127.0.0.1:9") + + def test_worker_with_unannotated_method_is_rejected(self): + """A method missing annotations is refused on the client too, matching the server.""" + + class Unannotated: + def demo_unannotated(self, x): + return x + + with pytest.raises(TypeError): + RpcWorkerHandle(Unannotated, server_url="http://127.0.0.1:9") + + async def test_trailing_slash_in_server_url(self, server, make_handle): + """A server url with a trailing slash still produces valid request paths.""" + handle = make_handle(f"{server.url}/") + assert await handle.demo_sync(a=1, b=1) == 2 + + async def test_client_and_server_agree_on_the_method_set(self, server): + """Every method the client exposes is routable on the server, and private ones are not.""" + handle = RpcWorkerHandle(E2eWorker, server_url=server.url) + assert "demo_sync" in handle._specs and "_bump" not in handle._specs + + async with httpx.AsyncClient(base_url=server.url, timeout=30.0, trust_env=False) as client: + for name in handle._specs: + response = await client.post( + f"/v1/{name}", json={"call_id": uuid.uuid4().hex, "query": {"__unknown__": 1}} + ) + assert response.status_code == 400, name + + private = await client.post("/v1/_bump", json={"call_id": uuid.uuid4().hex, "query": {}}) + assert private.status_code == 404 diff --git a/tests/fast/utils/workers/e2e/test_complex_types.py b/tests/fast/utils/workers/e2e/test_complex_types.py new file mode 100644 index 00000000000..5ecf8a04ec7 --- /dev/null +++ b/tests/fast/utils/workers/e2e/test_complex_types.py @@ -0,0 +1,188 @@ +import datetime +import uuid +from decimal import Decimal + +import pytest +from pydantic import ValidationError + +from tests.fast.utils.workers.e2e.e2e_worker import Colour, Item, Nested, Point + + +class TestComplexArgumentsAndResults: + async def test_nested_model_roundtrips(self, handle): + """A model containing models, dicts and sets survives the real wire both ways.""" + payload = Nested( + item=Item(name="a", values=[1, 2]), + lookup={"k": Item(name="b", values=[])}, + tags={"x", "y"}, + ) + assert await handle.demo_nested_model(payload=payload) == payload + + async def test_nested_model_result_is_revived_as_models(self, handle): + """The nested result arrives as model instances, not plain dicts.""" + payload = Nested(item=Item(name="a", values=[1]), lookup={"k": Item(name="b", values=[2])}, tags=set()) + result = await handle.demo_nested_model(payload=payload) + assert isinstance(result, Nested) + assert isinstance(result.item, Item) + assert isinstance(result.lookup["k"], Item) + + async def test_set_field_is_revived_as_a_set(self, handle): + """A set field arrives as a set even though json carries a list.""" + payload = Nested(item=Item(name="a", values=[]), lookup={}, tags={"x", "y"}) + result = await handle.demo_nested_model(payload=payload) + assert result.tags == {"x", "y"} + assert isinstance(result.tags, set) + + async def test_enum_roundtrips_as_the_member(self, handle): + """An enum argument and result stay enum members across the wire.""" + result = await handle.demo_enum(colour=Colour.BLUE) + assert result is Colour.BLUE + + async def test_dataclass_roundtrips_as_the_dataclass(self, handle): + """A dataclass argument is revived worker-side and the result comes back as a dataclass.""" + result = await handle.demo_dataclass(point=Point(x=1, y=2)) + assert isinstance(result, Point) + assert (result.x, result.y) == (2, 1) + + async def test_aware_datetime_keeps_its_offset(self, handle): + """A timezone-aware datetime keeps its instant and offset across the wire.""" + value = datetime.datetime(2026, 7, 27, 12, 30, tzinfo=datetime.timezone.utc) + result = await handle.demo_datetime(when=value) + assert isinstance(result, datetime.datetime) + assert result == value + + async def test_naive_datetime_stays_naive(self, handle): + """A naive datetime does not acquire a timezone on the way through.""" + value = datetime.datetime(2026, 7, 27, 12, 30) + result = await handle.demo_datetime(when=value) + assert result == value + assert result.tzinfo is None + + async def test_uuid_roundtrips_as_a_uuid(self, handle): + """A uuid argument and result stay uuid objects, not strings.""" + value = uuid.uuid4() + result = await handle.demo_uuid(value=value) + assert isinstance(result, uuid.UUID) + assert result == value + + async def test_decimal_keeps_precision(self, handle): + """A decimal keeps precision a float would lose.""" + value = Decimal("0.1234567890123456789") + assert await handle.demo_decimal(value=value) == value + + async def test_tuple_roundtrips_as_a_tuple(self, handle): + """A tuple result comes back as a tuple even though json carries a list.""" + result = await handle.demo_tuple(pair=(1, "a")) + assert isinstance(result, tuple) + assert result == (1, "a") + + @pytest.mark.parametrize("blob", [b"raw-bytes", b"", b"\x00\x80\xff", bytes(range(256))]) + async def test_bytes_roundtrip(self, handle, blob: bytes): + """A bytes argument and result stay bytes across a text protocol, including non-utf8 payloads.""" + assert await handle.demo_bytes(blob=blob) == blob + + async def test_bytes_inside_a_list_roundtrips(self, handle): + """Non-utf8 bytes nested in a container survive too.""" + blobs = [b"\x80", b"", b"\xff\xfe"] + assert await handle.demo_bytes_list(blobs=blobs) == blobs + + async def test_list_of_models_roundtrips(self, handle): + """A list of models arrives as a list of model instances.""" + items = [Item(name="a", values=[1]), Item(name="b", values=[2, 3])] + result = await handle.demo_model_list(items=items) + assert result == items + assert all(isinstance(item, Item) for item in result) + + async def test_empty_containers_roundtrip(self, handle): + """Empty containers survive rather than collapsing to null.""" + result = await handle.demo_model_list(items=[]) + assert result == [] + + async def test_optional_present_and_absent(self, handle): + """An optional argument roundtrips both with a value and with None.""" + assert await handle.demo_optional(value=7) == 7 + assert await handle.demo_optional(value=None) is None + + async def test_union_keeps_the_member_type(self, handle): + """A union result keeps the member type it was produced with.""" + assert await handle.demo_union(value=5) == 5 + assert await handle.demo_union(value="five") == "five" + + async def test_union_of_types_sharing_a_wire_form_resolves_to_str(self, handle): + """A datetime|str union shares one wire form, so the worker receives the str member.""" + when = datetime.datetime(2026, 7, 27, tzinfo=datetime.timezone.utc) + assert await handle.report_union_argument_type(value=when) == "str" + assert await handle.report_union_argument_type(value="plain") == "str" + + async def test_unicode_survives_the_wire(self, handle): + """Non-ascii text roundtrips byte for byte.""" + text = "δΈ­ζ–‡ πŸš€ \\ \" '" + result = await handle.demo_model(item=Item(name=text, values=[])) + assert result.name == text + + async def test_large_nested_payload_roundtrips(self, handle): + """A large nested payload roundtrips without truncation.""" + payload = Nested( + item=Item(name="big", values=list(range(5000))), + lookup={f"k{i}": Item(name=f"n{i}", values=[i]) for i in range(200)}, + tags={f"t{i}" for i in range(200)}, + ) + assert await handle.demo_nested_model(payload=payload) == payload + + async def test_none_result_roundtrips(self, handle): + """A method declared to return None reports success carrying None.""" + assert await handle.demo_none_result() is None + + +class TestWorkerSideTypes: + async def test_worker_receives_model_instances(self, handle): + """The worker is handed revived models, not the raw json dicts.""" + payload = Nested(item=Item(name="a", values=[1]), lookup={"k": Item(name="b", values=[])}, tags={"x"}) + assert await handle.report_nested_argument_types(payload=payload) == ["Nested", "Item", "Item", "set", "int"] + + async def test_worker_receives_a_dataclass_instance(self, handle): + """A dataclass argument reaches the worker as the dataclass.""" + assert await handle.report_dataclass_argument_type(point=Point(x=1, y=2)) == "Point" + + async def test_worker_receives_the_enum_member(self, handle): + """An enum argument reaches the worker as the member, comparable with is.""" + assert await handle.report_enum_argument_is_member(colour=Colour.BLUE) is True + + async def test_worker_receives_revived_scalars(self, handle): + """Scalars json cannot express reach the worker as their declared python types.""" + types = await handle.report_scalar_argument_types( + when=datetime.datetime(2026, 7, 27, tzinfo=datetime.timezone.utc), + value=uuid.uuid4(), + amount=Decimal("1.5"), + blob=b"x", + pair=(1, "a"), + ) + assert types == ["datetime", "UUID", "Decimal", "bytes", "tuple"] + + +class TestComplexTypeValidation: + async def test_wrong_model_type_is_rejected_before_sending(self, handle): + """A payload of the wrong shape is rejected client-side, not by the worker.""" + with pytest.raises(ValidationError): + await handle.demo_nested_model(payload={"item": {"name": "a"}}) + + async def test_unknown_enum_member_is_rejected(self, handle): + """A value outside the enum is refused rather than sent as a bare string.""" + with pytest.raises(ValidationError): + await handle.demo_enum(colour="green") + + async def test_wrong_tuple_arity_is_rejected(self, handle): + """A tuple of the wrong arity is refused client-side.""" + with pytest.raises(ValidationError): + await handle.demo_tuple(pair=(1, "a", 2)) + + async def test_extra_model_field_is_rejected(self, handle): + """An unknown field inside a model argument is refused.""" + with pytest.raises(ValidationError): + await handle.demo_model(item={"name": "a", "values": [], "extra": 1}) + + async def test_worker_stays_usable_after_a_rejected_call(self, handle): + """A client-side rejection does not consume anything server-side.""" + with pytest.raises(ValidationError): + await handle.demo_enum(colour="green") + assert await handle.demo_enum(colour=Colour.RED) is Colour.RED diff --git a/tests/fast/utils/workers/e2e/test_happy_path.py b/tests/fast/utils/workers/e2e/test_happy_path.py index b38016491b1..f53aa37b872 100644 --- a/tests/fast/utils/workers/e2e/test_happy_path.py +++ b/tests/fast/utils/workers/e2e/test_happy_path.py @@ -3,6 +3,7 @@ import httpx +from tests.fast.utils.workers.e2e.e2e_worker import Item from tests.fast.utils.workers.e2e.harness import READY_TIMEOUT_SECONDS @@ -16,6 +17,10 @@ async def test_async_method(self, handle): """An async method roundtrips a nested payload unchanged.""" assert await handle.demo_async(value={"k": [1, "x", None]}) == {"k": [1, "x", None]} + async def test_none_return(self, handle): + """A method declared -> None returns None rather than a missing value.""" + assert await handle.demo_none_result() is None + async def test_default_parameter_omitted(self, handle): """An omitted defaulted argument uses the worker-side default.""" assert await handle.demo_default_arg() == "hello world" @@ -75,6 +80,11 @@ async def test_concurrent_calls_from_several_handles(self, server, make_handle): class TestTypedPayloads: + async def test_pydantic_model_argument_and_result(self, handle): + """A pydantic model survives the wire in both directions as a real model.""" + result = await handle.demo_model(item=Item(name="x", values=[3, 1, 2])) + assert isinstance(result, Item) and result == Item(name="x", values=[3, 1, 2]) + async def test_scalar_types_keep_their_python_type(self, handle): """Scalars keep their type instead of collapsing to strings or ints.""" assert await handle.demo_async(value={"f": 1.5, "b": True, "s": "1", "n": None}) == { @@ -89,6 +99,16 @@ async def test_unicode_payload(self, handle): text = "δΈ­ζ–‡ πŸš€ \\ \" '" assert await handle.demo_async(value={"t": text}) == {"t": text} + async def test_large_upload(self, handle): + """A multi-megabyte request body is transferred intact.""" + blob = "x" * (4 * 1024 * 1024) + assert await handle.demo_large_upload(blob=blob) == len(blob) + + async def test_large_download(self, handle): + """A large response body is transferred intact.""" + values = await handle.demo_large_download(size=200_000) + assert len(values) == 200_000 and values[-1] == 199_999 + class TestManualProtocol: async def test_submit_then_poll_by_hand(self, raw): diff --git a/tests/fast/utils/workers/rpc/common/postponed_annotation_worker.py b/tests/fast/utils/workers/rpc/common/postponed_annotation_worker.py new file mode 100644 index 00000000000..5772eae844e --- /dev/null +++ b/tests/fast/utils/workers/rpc/common/postponed_annotation_worker.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from miles.utils.pydantic_utils import StrictBaseModel + + +class LatePayload(StrictBaseModel): + text: str + + +class PostponedWorker: + def demo_transform(self, payload: LatePayload) -> LatePayload: + return payload diff --git a/tests/fast/utils/workers/rpc/common/test_metadata.py b/tests/fast/utils/workers/rpc/common/test_metadata.py index 66545925c09..560d8680057 100644 --- a/tests/fast/utils/workers/rpc/common/test_metadata.py +++ b/tests/fast/utils/workers/rpc/common/test_metadata.py @@ -1,7 +1,11 @@ -from typing import Any +import functools +from collections.abc import Callable +from typing import Annotated, Any import pytest -from pydantic import ValidationError +from pydantic import Field, ValidationError + +from tests.fast.utils.workers.rpc.common.postponed_annotation_worker import LatePayload, PostponedWorker from miles.utils.pydantic_utils import StrictBaseModel from miles.utils.workers.rpc.common.metadata import DEFAULT_CONCURRENCY_GROUP, collect_rpc_method_specs, rpc @@ -12,6 +16,22 @@ class _Payload(StrictBaseModel): count: int = 1 +def _passthrough(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return fn(*args, **kwargs) + + return wrapper + + +def _opaque_passthrough(fn: Callable[..., Any]) -> Callable[..., Any]: + def wrapper(*args: Any, **kwargs: Any) -> Any: + return fn(*args, **kwargs) + + wrapper.__wrapped__ = fn + return wrapper + + class _GoodWorker: demo_class_attribute = 3 @@ -67,53 +87,348 @@ def test_is_async_flag(self): specs = collect_rpc_method_specs(_GoodWorker) assert specs["demo_async_model"].is_async and not specs["demo_default_arg"].is_async - def test_async_method_with_non_default_concurrency_group_rejected(self): - """A concurrency group on an async method would be silently ignored, so it is refused.""" + +class TestDecoratorChainConcurrencyGroup: + def test_marker_above_wrapper_is_found(self): + """@rpc applied outside a functools.wraps wrapper still declares its concurrency group.""" class Worker: @rpc(concurrency_group="heavy") - async def demo_async_grouped(self) -> None: - pass + @_passthrough + def demo_marker_outside(self, x: int) -> int: + return x - with pytest.raises(TypeError, match="concurrency groups only serialize sync methods"): - collect_rpc_method_specs(Worker) + specs = collect_rpc_method_specs(Worker) + assert specs["demo_marker_outside"].concurrency_group == "heavy" + + def test_marker_below_wrapper_is_found(self): + """@rpc applied inside a functools.wraps wrapper still declares its concurrency group.""" + + class Worker: + @_passthrough + @rpc(concurrency_group="heavy") + def demo_marker_inside(self, x: int) -> int: + return x + + specs = collect_rpc_method_specs(Worker) + assert specs["demo_marker_inside"].concurrency_group == "heavy" + + def test_marker_between_two_wrapper_layers_is_found(self): + """@rpc sandwiched between two wrapper layers is still found by walking the decorator chain.""" + + class Worker: + @_passthrough + @rpc(concurrency_group="heavy") + @_passthrough + def demo_marker_nested(self, x: int) -> int: + return x + + specs = collect_rpc_method_specs(Worker) + assert specs["demo_marker_nested"].concurrency_group == "heavy" + + def test_marker_hidden_by_a_wrapper_that_copies_nothing_is_found(self): + """A wrapper that does not copy the wrapped function's attributes cannot hide the marker.""" + + class Worker: + @_opaque_passthrough + @rpc(concurrency_group="heavy") + def demo_marker_hidden(self, x: int) -> int: + return x + + specs = collect_rpc_method_specs(Worker) + assert specs["demo_marker_hidden"].concurrency_group == "heavy" + + def test_outermost_marker_wins_over_an_inner_one(self): + """When two decorator layers each declare a group, the outermost declaration decides.""" + + class Worker: + @rpc(concurrency_group="outer") + @_opaque_passthrough + @rpc(concurrency_group="inner") + def demo_marker_conflict(self, x: int) -> int: + return x + + specs = collect_rpc_method_specs(Worker) + assert specs["demo_marker_conflict"].concurrency_group == "outer" class TestQueryModel: def test_decode_query_applies_defaults(self): - """An omitted defaulted parameter is filled in on the server side.""" + """Omitted parameters with defaults resolve to their default values.""" specs = collect_rpc_method_specs(_GoodWorker) - assert specs["demo_default_arg"].serializer.decode_query({"a": 1}) == {"a": 1, "b": 10} + assert specs["demo_default_arg"].serializer.decode_query({"a": 5}) == {"a": 5, "b": 10} def test_decode_query_parses_nested_model(self): - """A nested model argument is revived as a model instance.""" + """Nested pydantic payloads are revived into real model instances.""" specs = collect_rpc_method_specs(_GoodWorker) - decoded = specs["demo_async_model"].serializer.decode_query({"payload": {"text": "hi"}}) - assert decoded == {"payload": _Payload(text="hi")} + kwargs = specs["demo_async_model"].serializer.decode_query({"payload": {"text": "hi"}}) + assert kwargs["payload"] == _Payload(text="hi") def test_missing_required_param_rejected(self): - """A missing required argument fails validation instead of defaulting.""" + """Missing required parameters raise a validation error.""" specs = collect_rpc_method_specs(_GoodWorker) with pytest.raises(ValidationError): specs["demo_default_arg"].serializer.decode_query({}) def test_unknown_param_rejected(self): - """An argument the method does not declare fails validation.""" + """Extra unknown parameters raise a validation error.""" specs = collect_rpc_method_specs(_GoodWorker) with pytest.raises(ValidationError): - specs["demo_default_arg"].serializer.decode_query({"a": 1, "nope": 2}) + specs["demo_default_arg"].serializer.decode_query({"a": 1, "unknown": 2}) def test_wrong_type_rejected(self): - """An argument that cannot be coerced fails validation.""" + """Type-mismatched parameters raise a validation error.""" specs = collect_rpc_method_specs(_GoodWorker) with pytest.raises(ValidationError): - specs["demo_default_arg"].serializer.decode_query({"a": "not-int"}) + specs["demo_default_arg"].serializer.decode_query({"a": "not-an-int"}) + + +class TestParameterKinds: + def test_keyword_only_parameters_are_supported(self): + """Keyword-only parameters are accepted, land in the query model and decode like normal ones.""" + + class Worker: + def demo_keyword_only(self, *, required: int, optional: str = "fallback") -> int: + return required + + specs = collect_rpc_method_specs(Worker) + serializer = specs["demo_keyword_only"].serializer + assert serializer.decode_query({"required": 5}) == {"required": 5, "optional": "fallback"} + with pytest.raises(ValidationError): + serializer.decode_query({"optional": "only"}) + + def test_positional_only_receiver_is_supported(self): + """A positional-only self is accepted because the receiver never reaches the wire.""" + + class Worker: + def demo_positional_receiver(self, /, value: int) -> int: + return value + + specs = collect_rpc_method_specs(Worker) + assert specs["demo_positional_receiver"].serializer.decode_query({"value": 7}) == {"value": 7} + + +class TestAnnotatedParameters: + def test_annotated_parameter_constraints_are_preserved(self): + """Constraints carried by Annotated metadata survive hint resolution and are enforced.""" + + class Worker: + def demo_constrained(self, value: Annotated[int, Field(ge=1)]) -> int: + return value + + serializer = collect_rpc_method_specs(Worker)["demo_constrained"].serializer + assert serializer.decode_query({"value": 3}) == {"value": 3} + with pytest.raises(ValidationError): + serializer.decode_query({"value": 0}) + + +class TestPostponedAnnotations: + def test_string_annotations_resolved_in_worker_module(self): + """A worker module using postponed annotations still builds real typed models.""" + specs = collect_rpc_method_specs(PostponedWorker) + kwargs = specs["demo_transform"].serializer.decode_query({"payload": {"text": "hi"}}) + assert kwargs["payload"] == LatePayload(text="hi") + + def test_string_return_annotation_resolved(self): + """A postponed return annotation resolves into a working result adapter.""" + specs = collect_rpc_method_specs(PostponedWorker) + assert specs["demo_transform"].serializer.decode_result({"text": "hi"}) == LatePayload(text="hi") + + +class TestInheritance: + def test_inherited_methods_collected(self): + """Methods inherited from a base worker class are exposed too.""" + + class Child(_GoodWorker): + def demo_child_only(self, x: int) -> int: + return x + + specs = collect_rpc_method_specs(Child) + assert {"demo_default_arg", "demo_async_model", "demo_grouped", "demo_child_only"} <= set(specs) class TestResultAdapter: def test_result_roundtrip(self): - """A model result encodes to json-safe data and decodes back to a model.""" - serializer = collect_rpc_method_specs(_GoodWorker)["demo_async_model"].serializer - encoded: Any = serializer.encode_result(_Payload(text="hi", count=2)) - assert encoded == {"text": "hi", "count": 2} - assert serializer.decode_result(encoded) == _Payload(text="hi", count=2) + """Return values are encoded as plain json data and decode back into the model.""" + specs = collect_rpc_method_specs(_GoodWorker) + serializer = specs["demo_async_model"].serializer + dumped = serializer.encode_result(_Payload(text="hi")) + assert not isinstance(dumped, _Payload) + assert dumped == {"text": "hi", "count": 1} + assert serializer.decode_result(dumped) == _Payload(text="hi") + + def test_none_return_annotation(self): + """Methods annotated -> None get a NoneType result adapter.""" + specs = collect_rpc_method_specs(_GoodWorker) + assert specs["demo_grouped"].serializer.decode_result(None) is None + + +class TestFailLoud: + def test_async_method_with_non_default_concurrency_group_rejected(self): + """An async method with a non-default concurrency group fails at collection time.""" + + class Worker: + @rpc(concurrency_group="train") + async def demo_async_grouped(self) -> int: + return 0 + + with pytest.raises(TypeError, match="concurrency_group"): + collect_rpc_method_specs(Worker) + + def test_missing_param_annotation_rejected(self): + """A parameter without a type annotation fails at collection time.""" + + class Worker: + def demo_unannotated_arg(self, x) -> int: + return 0 + + with pytest.raises(TypeError, match="must be type-annotated"): + collect_rpc_method_specs(Worker) + + def test_missing_return_annotation_rejected(self): + """A method without a return annotation fails at collection time.""" + + class Worker: + def demo_unannotated_return(self, x: int): + return x + + with pytest.raises(TypeError, match="return type annotation"): + collect_rpc_method_specs(Worker) + + def test_var_positional_rejected(self): + """*args signatures fail at collection time.""" + + class Worker: + def demo_var_positional(self, *x: int) -> int: + return 0 + + with pytest.raises(TypeError, match="args"): + collect_rpc_method_specs(Worker) + + def test_var_keyword_rejected(self): + """**kwargs signatures fail at collection time.""" + + class Worker: + def demo_var_keyword(self, **x: int) -> int: + return 0 + + with pytest.raises(TypeError, match="kwargs"): + collect_rpc_method_specs(Worker) + + def test_positional_only_rejected(self): + """Positional-only parameters fail at collection time since calls pass kwargs.""" + + class Worker: + def demo_positional_only(self, x: int, /) -> int: + return x + + with pytest.raises(TypeError, match="positional-only"): + collect_rpc_method_specs(Worker) + + def test_non_self_receiver_rejected(self): + """An unconventionally named receiver is refused rather than silently dropped.""" + + class Worker: + def demo_odd_receiver(this, x: int) -> int: + return x + + with pytest.raises(TypeError, match="receiver parameter 'self'"): + collect_rpc_method_specs(Worker) + + def test_forgotten_self_is_rejected_instead_of_eating_the_first_argument(self): + """A method that forgets self would otherwise lose its first parameter off the wire.""" + + class Worker: + def demo_forgot_self(a: int, b: int) -> int: + return a + b + + with pytest.raises(TypeError, match="receiver parameter 'self'"): + collect_rpc_method_specs(Worker) + + def test_keyword_only_receiver_rejected(self): + """A keyword-only self is refused at collection time instead of blowing up on call.""" + + class Worker: + def demo_keyword_only_receiver(*, self) -> int: + return 0 + + with pytest.raises(TypeError, match="receiver parameter positionally"): + collect_rpc_method_specs(Worker) + + def test_method_without_any_parameter_rejected(self): + """A method taking no parameters at all is refused for lacking a receiver.""" + + class Worker: + def demo_no_parameters() -> int: + return 0 + + with pytest.raises(TypeError, match="must take a receiver parameter"): + collect_rpc_method_specs(Worker) + + def test_public_nested_model_class_rejected_as_non_method(self): + """A public nested model class is refused as not being a method at all.""" + + class Worker: + class Config(StrictBaseModel): + text: str + + def demo_ok(self, x: int) -> int: + return x + + with pytest.raises(TypeError) as excinfo: + collect_rpc_method_specs(Worker) + assert "not a method" in str(excinfo.value) + assert "receiver" not in str(excinfo.value) + + def test_public_class_alias_attribute_rejected_as_non_method(self): + """A public class alias attribute is refused as not being a method at all.""" + + class Worker: + demo_alias = _Payload + + def demo_ok(self, x: int) -> int: + return x + + with pytest.raises(TypeError) as excinfo: + collect_rpc_method_specs(Worker) + assert "not a method" in str(excinfo.value) + assert "receiver" not in str(excinfo.value) + + def test_wrapped_async_method_stays_async(self): + """A functools.wraps-decorated async method is still detected as async.""" + + def passthrough(fn): + @functools.wraps(fn) + async def wrapper(*args, **kwargs): + return await fn(*args, **kwargs) + + return wrapper + + class Worker: + @passthrough + async def demo_wrapped_async(self, x: int) -> int: + return x + + specs = collect_rpc_method_specs(Worker) + assert specs["demo_wrapped_async"].is_async + assert specs["demo_wrapped_async"].serializer.decode_query({"x": 1}) == {"x": 1} + + def test_no_public_methods_rejected(self): + """A worker class with no public methods fails at collection time.""" + + class Worker: + def _demo_hidden(self, x: int) -> int: + return x + + with pytest.raises(TypeError, match="no public rpc methods"): + collect_rpc_method_specs(Worker) + + def test_any_annotation_allowed(self): + """Any-annotated parameters are accepted and passed through.""" + + class Worker: + def demo_any(self, x: Any) -> Any: + return x + + specs = collect_rpc_method_specs(Worker) + assert specs["demo_any"].serializer.decode_query({"x": [1, "a"]}) == {"x": [1, "a"]} diff --git a/tests/fast/utils/workers/rpc/common/test_serialization.py b/tests/fast/utils/workers/rpc/common/test_serialization.py new file mode 100644 index 00000000000..1b28f6c0588 --- /dev/null +++ b/tests/fast/utils/workers/rpc/common/test_serialization.py @@ -0,0 +1,317 @@ +import dataclasses +import datetime +import enum +import json +import uuid +from decimal import Decimal +from pathlib import PurePosixPath +from typing import Any, Literal + +import pytest +from pydantic import ValidationError +from pydantic_core import PydanticSerializationError +from typing_extensions import TypedDict + +from miles.utils.pydantic_utils import StrictBaseModel +from miles.utils.workers.rpc.common.serialization import RpcSerializer + + +class Colour(enum.Enum): + RED = "red" + BLUE = "blue" + + +class Level(enum.IntEnum): + LOW = 1 + HIGH = 2 + + +class Inner(StrictBaseModel): + name: str + score: float + + +class Outer(StrictBaseModel): + inner: Inner + tags: list[str] + lookup: dict[str, Inner] + + +@dataclasses.dataclass +class Point: + x: int + y: int + + +class Blob(StrictBaseModel): + data: bytes + + +class Options(TypedDict): + retries: int + label: str + + +_ROUNDTRIP_CASES = [ + ("int", int, 42), + ("negative_int", int, -7), + ("big_int", int, 2**70), + ("float", float, 1.5), + ("bool", bool, True), + ("str", str, "text"), + ("unicode_str", str, "δΈ­ζ–‡ πŸš€ \\ \" '"), + ("empty_str", str, ""), + ("none", type(None), None), + ("list_of_int", list[int], [1, 2, 3]), + ("empty_list", list[int], []), + ("nested_list", list[list[int]], [[1], [2, 3], []]), + ("dict_str_int", dict[str, int], {"a": 1}), + ("dict_of_lists", dict[str, list[int]], {"a": [1, 2]}), + ("tuple", tuple[int, str], (1, "a")), + ("variadic_tuple", tuple[int, ...], (1, 2, 3)), + ("set", set[int], {1, 2, 3}), + ("frozenset", frozenset[str], frozenset({"a", "b"})), + ("optional_present", int | None, 5), + ("optional_absent", int | None, None), + ("union", int | str, "either"), + ("literal", Literal["a", "b"], "b"), + ("str_enum", Colour, Colour.RED), + ("int_enum", Level, Level.HIGH), + ("uuid", uuid.UUID, uuid.UUID("12345678-1234-5678-1234-567812345678")), + ("datetime", datetime.datetime, datetime.datetime(2026, 7, 27, 12, 30, tzinfo=datetime.timezone.utc)), + ("naive_datetime", datetime.datetime, datetime.datetime(2026, 7, 27, 12, 30)), + ("date", datetime.date, datetime.date(2026, 7, 27)), + ("time", datetime.time, datetime.time(12, 30, 15)), + ("timedelta", datetime.timedelta, datetime.timedelta(seconds=90)), + ("decimal", Decimal, Decimal("1.25")), + ("path", PurePosixPath, PurePosixPath("/tmp/x")), + ("bytes", bytes, b"raw-bytes"), + ("model", Inner, Inner(name="x", score=1.5)), + ( + "nested_model", + Outer, + Outer(inner=Inner(name="x", score=1.0), tags=["a"], lookup={"k": Inner(name="y", score=2.0)}), + ), + ("list_of_models", list[Inner], [Inner(name="a", score=1.0), Inner(name="b", score=2.0)]), + ("dict_of_models", dict[str, Inner], {"k": Inner(name="a", score=1.0)}), + ("optional_model", Inner | None, Inner(name="a", score=1.0)), + ("dataclass", Point, Point(x=1, y=2)), + ("typed_dict", Options, {"retries": 2, "label": "x"}), + ("any_scalar", Any, 3), + ("any_container", Any, {"k": [1, {"n": None}]}), + ("deeply_nested", dict[str, list[dict[str, int]]], {"a": [{"b": 1}, {"c": 2}]}), +] + + +_BYTES_CASES = [ + ("non_utf8_bytes", bytes, b"\x00\x80\xff"), + ("empty_bytes", bytes, b""), + ("list_of_bytes", list[bytes], [b"\x00\x80\xff", b""]), + ("dict_of_bytes", dict[str, bytes], {"k": b"\x00\x80\xff"}), +] + + +def _serializer(annotation: type) -> RpcSerializer: + return RpcSerializer.create( + query_model_name="Query", query_fields={"payload": (annotation, ...)}, result_annotation=annotation + ) + + +def _through_the_wire(payload: Any) -> Any: + return json.loads(json.dumps(payload)) + + +@pytest.mark.parametrize("case", _ROUNDTRIP_CASES, ids=[case[0] for case in _ROUNDTRIP_CASES]) +class TestWireRoundtrip: + def test_result_survives_the_wire_as_the_declared_type(self, case): + """Every supported result type comes back from json as the declared python type.""" + _, annotation, value = case + serializer = _serializer(annotation) + assert serializer.decode_result(_through_the_wire(serializer.encode_result(value))) == value + + def test_argument_survives_the_wire_as_the_declared_type(self, case): + """Every supported argument type reaches the worker as the declared python type.""" + _, annotation, value = case + serializer = _serializer(annotation) + assert ( + serializer.decode_query(_through_the_wire(serializer.encode_query({"payload": value})))["payload"] == value + ) + + +class TestTypeRevival: + def test_model_result_is_revived_as_a_model(self): + """A model result arrives as an instance rather than the raw json dict.""" + serializer = _serializer(Inner) + revived = serializer.decode_result(_through_the_wire(serializer.encode_result(Inner(name="x", score=1.0)))) + assert isinstance(revived, Inner) + + def test_nested_model_argument_is_revived_all_the_way_down(self): + """A nested model argument arrives with its inner models revived too.""" + value = Outer(inner=Inner(name="x", score=1.0), tags=[], lookup={"k": Inner(name="y", score=2.0)}) + serializer = _serializer(Outer) + revived = serializer.decode_query(_through_the_wire(serializer.encode_query({"payload": value})))["payload"] + assert isinstance(revived.lookup["k"], Inner) + + def test_dataclass_argument_is_revived_as_the_dataclass(self): + """A dataclass argument is handed to the method as an instance, not a dict.""" + serializer = _serializer(Point) + revived = serializer.decode_query(_through_the_wire(serializer.encode_query({"payload": Point(x=1, y=2)}))) + assert isinstance(revived["payload"], Point) + + def test_model_argument_is_revived_as_a_model(self): + """A model argument is handed to the method as an instance, not a dict.""" + serializer = _serializer(Inner) + revived = serializer.decode_query( + _through_the_wire(serializer.encode_query({"payload": Inner(name="x", score=1.0)})) + ) + assert isinstance(revived["payload"], Inner) + + def test_enum_result_is_revived_as_the_enum_member(self): + """An enum result arrives as the member rather than its bare value.""" + serializer = _serializer(Colour) + assert serializer.decode_result(_through_the_wire(serializer.encode_result(Colour.BLUE))) is Colour.BLUE + + def test_int_enum_result_is_revived_as_the_enum_member_not_a_bare_int(self): + """An int enum result arrives as the member itself rather than the equal plain int.""" + serializer = _serializer(Level) + revived = serializer.decode_result(_through_the_wire(serializer.encode_result(Level.HIGH))) + assert type(revived) is Level + assert revived is Level.HIGH + + def test_int_enum_argument_is_revived_as_the_enum_member_not_a_bare_int(self): + """An int enum argument reaches the method as the member itself rather than the equal plain int.""" + serializer = _serializer(Level) + revived = serializer.decode_query(_through_the_wire(serializer.encode_query({"payload": Level.LOW}))) + assert type(revived["payload"]) is Level + assert revived["payload"] is Level.LOW + + def test_tuple_result_is_revived_as_a_tuple(self): + """A tuple result arrives as a tuple even though json carries a list.""" + serializer = _serializer(tuple[int, str]) + assert isinstance(serializer.decode_result(_through_the_wire(serializer.encode_result((1, "a")))), tuple) + + def test_set_result_is_revived_as_a_set(self): + """A set result arrives as a set even though json carries a list.""" + serializer = _serializer(set[int]) + assert isinstance(serializer.decode_result(_through_the_wire(serializer.encode_result({1, 2}))), set) + + def test_datetime_result_is_revived_as_a_datetime(self): + """A datetime result arrives as a datetime rather than an iso string.""" + serializer = _serializer(datetime.datetime) + value = datetime.datetime(2026, 7, 27, 12, 0, tzinfo=datetime.timezone.utc) + revived = serializer.decode_result(_through_the_wire(serializer.encode_result(value))) + assert isinstance(revived, datetime.datetime) + assert revived.utcoffset() == value.utcoffset() + + def test_dataclass_result_is_revived_as_the_dataclass(self): + """A dataclass result arrives as an instance rather than a dict.""" + serializer = _serializer(Point) + assert isinstance( + serializer.decode_result(_through_the_wire(serializer.encode_result(Point(x=1, y=2)))), Point + ) + + def test_decimal_result_keeps_its_precision(self): + """A decimal result keeps full precision instead of degrading to a float.""" + serializer = _serializer(Decimal) + value = Decimal("0.1234567890123456789") + assert serializer.decode_result(_through_the_wire(serializer.encode_result(value))) == value + + def test_int_result_does_not_widen_to_float(self): + """An int result stays an int across the wire.""" + serializer = _serializer(int) + assert isinstance(serializer.decode_result(_through_the_wire(serializer.encode_result(1))), int) + + def test_bool_result_does_not_collapse_to_int(self): + """A bool result stays a bool rather than becoming 0 or 1.""" + serializer = _serializer(bool) + assert serializer.decode_result(_through_the_wire(serializer.encode_result(True))) is True + + def test_encoded_payloads_are_plain_json_types(self): + """Encoding produces json-native values so the wire never sees python objects.""" + serializer = _serializer(Outer) + value = Outer(inner=Inner(name="x", score=1.0), tags=["a"], lookup={}) + assert json.dumps(serializer.encode_result(value)) + + +@pytest.mark.parametrize("case", _BYTES_CASES, ids=[case[0] for case in _BYTES_CASES]) +class TestBytesOnTheWire: + def test_bytes_result_survives_the_wire(self, case): + """Arbitrary bytes results come back byte for byte after base64 transport.""" + _, annotation, value = case + serializer = _serializer(annotation) + assert serializer.decode_result(_through_the_wire(serializer.encode_result(value))) == value + + def test_bytes_argument_survives_the_wire(self, case): + """Arbitrary bytes arguments reach the worker byte for byte after base64 transport.""" + _, annotation, value = case + serializer = _serializer(annotation) + decoded = serializer.decode_query(_through_the_wire(serializer.encode_query({"payload": value}))) + assert decoded["payload"] == value + + def test_encoded_bytes_are_json_encodable(self, case): + """Encoding bytes yields a json encodable payload rather than raw python bytes.""" + _, annotation, value = case + serializer = _serializer(annotation) + assert json.dumps(serializer.encode_result(value), allow_nan=False) + + +class TestBytesEncoding: + def test_encoded_non_utf8_bytes_result_is_a_string(self): + """A non-utf8 bytes result encodes to a base64 string on the wire.""" + serializer = _serializer(bytes) + assert isinstance(serializer.encode_result(b"\x00\x80\xff"), str) + + def test_encoded_non_utf8_bytes_argument_is_a_string(self): + """A non-utf8 bytes argument encodes to a base64 string on the wire.""" + serializer = _serializer(bytes) + assert isinstance(serializer.encode_query({"payload": b"\x00\x80\xff"})["payload"], str) + + def test_utf8_bytes_nested_in_a_model_survive_the_wire(self): + """Utf8-decodable bytes stored in a model field come back byte for byte on the revived model.""" + serializer = _serializer(Blob) + decoded = serializer.decode_result(_through_the_wire(serializer.encode_result(Blob(data=b"payload")))) + assert decoded.data == b"payload" + + def test_non_utf8_bytes_nested_in_a_model_are_refused_loudly(self): + """A model carries its own serialization config, so non-utf8 bytes in a model field fail instead of corrupting.""" + serializer = _serializer(Blob) + with pytest.raises(UnicodeDecodeError): + serializer.encode_result(Blob(data=b"\x00\x80\xff")) + + +class TestRejectedPayloads: + def test_unserializable_result_is_rejected(self): + """A result that is not json encodable fails rather than being silently coerced.""" + serializer = _serializer(int) + with pytest.raises((TypeError, PydanticSerializationError)): + json.dumps(serializer.encode_result(object())) + + def test_wrong_result_type_is_rejected(self): + """A result payload that does not match the annotation fails validation.""" + serializer = _serializer(int) + with pytest.raises(ValidationError): + serializer.decode_result("not-an-int") + + def test_extra_model_field_is_rejected(self): + """Strict models refuse unknown fields arriving from the wire.""" + serializer = _serializer(Inner) + with pytest.raises(ValidationError): + serializer.decode_result({"name": "x", "score": 1.0, "extra": 1}) + + def test_missing_model_field_is_rejected(self): + """Strict models refuse payloads missing declared fields.""" + serializer = _serializer(Inner) + with pytest.raises(ValidationError): + serializer.decode_result({"name": "x"}) + + def test_unknown_argument_is_rejected(self): + """An argument the method does not declare is refused rather than ignored.""" + serializer = _serializer(int) + with pytest.raises(ValidationError): + serializer.decode_query({"payload": 1, "unknown": 2}) + + def test_wrong_argument_type_is_rejected(self): + """An argument that does not match the annotation fails validation.""" + serializer = _serializer(Inner) + with pytest.raises(ValidationError): + serializer.decode_query({"payload": "not-a-model"})