From 16904b98c7ab920ed1fb81bd3547c11e8e4e69ff Mon Sep 17 00:00:00 2001 From: Nicholas Devenish Date: Mon, 26 Oct 2026 14:03:00 +0000 Subject: [PATCH] Add typed_handler decorator to validate recipe messages Service callbacks can annotate their message, and an optional parameters argument, with pydantic models. By default the parameters come from the recipe step only; parameters_from lets a service merge in values from the message, or nested dictionaries such as message.parameters, in an explicit order of precedence. Services can opt in to substituting values from the recipe environment. Messages failing validation are logged and nacked without requeueing, or optionally dropped. CommonService.subscribe now rejects callbacks that can not be called with (rw, header, message), and points callbacks declaring a parameters argument at typed_handler, instead of failing when the first message arrives. Also adds a NonEmpty validator, for fields such as Annotated[Path, NonEmpty]. pydantic is added as a dependency. --- .pre-commit-config.yaml | 2 +- pyproject.toml | 3 +- src/workflows/recipe/__init__.py | 5 + src/workflows/recipe/typed.py | 317 ++++++++++++++++++++++ src/workflows/services/common_service.py | 32 +++ tests/recipe/test_typed.py | 319 +++++++++++++++++++++++ tests/services/test_common_service.py | 37 +++ 7 files changed, 713 insertions(+), 2 deletions(-) create mode 100644 src/workflows/recipe/typed.py create mode 100644 tests/recipe/test_typed.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 34237922..53dfd10b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,5 +36,5 @@ repos: hooks: - id: mypy files: 'src/.*\.py$' - additional_dependencies: ['types-setuptools>=57.0.2'] + additional_dependencies: ['types-setuptools>=57.0.2', 'pydantic>=2'] language_version: python3.12 diff --git a/pyproject.toml b/pyproject.toml index 5854fd53..22b979c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,8 @@ dependencies = [ "stomp-py>=7", "opentelemetry-api", "opentelemetry-sdk", - "opentelemetry-exporter-otlp-proto-http" + "opentelemetry-exporter-otlp-proto-http", + "pydantic>=2", ] [project.urls] diff --git a/src/workflows/recipe/__init__.py b/src/workflows/recipe/__init__.py index 0f7fc84f..f8bbb892 100644 --- a/src/workflows/recipe/__init__.py +++ b/src/workflows/recipe/__init__.py @@ -23,6 +23,11 @@ logger = logging.getLogger("workflows.recipe") +# Attribute set on callbacks decorated with workflows.recipe.typed.typed_handler, +# which accept recipe messages with the usual three arguments despite their +# signature declaring an additional 'parameters' argument +TYPED_HANDLER_MARKER = "_workflows_typed_handler" + def _wrap_subscription( transport_layer: CommonTransport, diff --git a/src/workflows/recipe/typed.py b/src/workflows/recipe/typed.py new file mode 100644 index 00000000..2f9c36a4 --- /dev/null +++ b/src/workflows/recipe/typed.py @@ -0,0 +1,317 @@ +"""Typed message handlers for workflows services. + +A service callback decorated with :func:`typed_handler` can annotate its +``message`` argument, and optionally declare a ``parameters`` argument, with +pydantic models (or any type a :class:`pydantic.TypeAdapter` understands). +Incoming messages are validated before the callback is invoked, and messages +failing validation are logged and rejected:: + + class Parameters(pydantic.BaseModel): + output_filename: Annotated[Path, NonEmpty] + timeout: int = 60 + + @typed_handler + def receive(self, rw, header, message: dict[str, Any], parameters: Parameters): + ... + +Parameter sources +----------------- + +By default ``parameters`` is built only from the ``parameters`` dictionary of +the current recipe step. Services that also let incoming messages supply or +override parameters say so with ``parameters_from``, an ordered sequence of +sources with the **highest precedence first**: + +``"step"`` + The ``parameters`` dictionary of the current recipe step. +``"message"`` + The top level of the message payload. Every key of the message becomes + a candidate parameter, including the message's actual data, so this is + best used with a parameters model that ignores unknown fields (the + pydantic default). +``"step."`` / ``"message."`` + A dictionary nested inside one of the above, eg. ``"message.parameters"`` + or ``"step.ispyb_parameters"``. Further levels may be chained with dots, + eg. ``"message.parameters.extra"``. + +The sources are merged as follows: + +* Merging is shallow: a top-level key from a higher precedence source + replaces the whole value from lower ones. Nested dictionaries are not + combined. +* A key that is present takes precedence even when its value is falsy, eg. + ``0``, ``""``, ``False`` or ``None``. This differs from the + ``message.get(key) or default`` lookups used by some older services. +* A source that does not exist, or is not a dictionary, is skipped. For + example ``"step"`` contributes nothing to a message that is not part of a + recipe, and ``"message"`` contributes nothing if the payload is a list. + +Variable substitution +--------------------- + +With ``substitute=True``, after merging, ``$name`` and ``${name}`` in string +values (including inside nested lists and dictionaries) are replaced with +values from the recipe environment. Unknown names are left in place, and +``$$`` becomes ``$``. + +Substitution is off by default because parameters often contain text where +``$`` means something else, such as shell commands (``$$``, ``$HOME``) or +email content. A name that happens to match a recipe environment key, eg. +``$ID``, would otherwise be silently replaced. + +Some common configurations: + +``("step",)`` (the default) + Recipe step parameters only. The message can not influence them. +``("message", "step")`` + Any top-level message key overrides the recipe step parameter of the + same name. +``("message.parameters", "step")`` + Only keys in the message's ``parameters`` dictionary override recipe + step parameters. +``("message.parameters", "step.ispyb_parameters", "step")`` + As above, with the step's ``ispyb_parameters`` also overriding the rest + of the step parameters. +``("step.ispyb_parameters", "message", "step")`` + The step's ``ispyb_parameters`` take precedence even over the message. + +Only let messages override parameters when senders are meant to control +the service's behaviour. In a service whose message is data to be stored or +forwarded, a key in that data could otherwise silently change how it is +handled. +""" + +from __future__ import annotations + +import functools +import inspect +import logging +import string +import typing +from collections.abc import Callable, Mapping, Sequence +from typing import Any, Literal + +import pydantic + +import workflows.recipe + +logger = logging.getLogger("workflows.recipe.typed") + +# Annotations that mean an argument is passed through without validation: +# explicitly typed as Any, or not annotated at all (inspect.Parameter.empty is +# the sentinel used when get_type_hints has no entry for the argument) +_UNVALIDATED = (Any, inspect.Parameter.empty) + + +def _require_nonempty(value: Any) -> Any: + if value == "": + raise ValueError("must not be empty") + return value + + +# Validator rejecting empty strings before conversion, eg. Annotated[Path, NonEmpty]. +# Without it Path("") silently becomes Path("."), and constraints such as +# min_length can not be applied to a Path. +NonEmpty = pydantic.BeforeValidator(_require_nonempty) + + +def _substitute(value: Any, substitutions: Mapping[str, Any]) -> Any: + """Recursively replace $variables in strings with values from substitutions""" + if isinstance(value, str) and "$" in value: + return string.Template(value).safe_substitute(substitutions) + if isinstance(value, Mapping): + return {k: _substitute(v, substitutions) for k, v in value.items()} + if isinstance(value, list): + return [_substitute(v, substitutions) for v in value] + return value + + +_DEFAULT_PARAMETERS_FROM = ("step",) + + +def _parse_sources(sources: Sequence[str]) -> list[tuple[str, list[str]]]: + """Split parameter source strings into (root, [nested keys]) pairs""" + if isinstance(sources, str): + raise TypeError( + f"parameters_from must be a sequence of sources, not a string: " + f"use ({sources!r},)" + ) + if not sources: + raise ValueError("parameters_from must name at least one source") + parsed = [] + for source in sources: + root, *keys = source.split(".") + if root not in ("step", "message") or "" in keys: + raise ValueError( + f"Invalid parameter source {source!r}: sources must be 'step' or " + "'message', optionally followed by '.key' to select a nested dictionary" + ) + parsed.append((root, keys)) + return parsed + + +def merged_parameters( + rw: workflows.recipe.RecipeWrapper | None, + message: Any, + *, + parameters_from: Sequence[str] = _DEFAULT_PARAMETERS_FROM, + substitute: bool = False, +) -> dict[str, Any]: + """Assemble the effective parameters for a recipe step from several sources. + + See the module documentation for a full description of the rules. + + Args: + rw: The recipe wrapper of the incoming message, or None for messages + that are not part of a recipe. + message: The incoming message payload. + parameters_from: The sources to merge, highest precedence first. Each + is 'step' (the recipe step parameters) or 'message' (the message + payload), optionally followed by '.key' to select a nested + dictionary, eg. 'message.parameters' or 'step.ispyb_parameters'. + substitute: Whether to replace $variables in string values with + values from the recipe environment. + + Returns: + A new dictionary of the merged parameters. + """ + roots: dict[str, Any] = { + "step": rw.recipe_step.get("parameters") if rw and rw.recipe_step else None, + "message": message, + } + parameters: dict[str, Any] = {} + # Apply lowest precedence first, so higher precedence sources overwrite + for root, keys in reversed(_parse_sources(parameters_from)): + layer = roots[root] + for key in keys: + layer = layer.get(key) if isinstance(layer, Mapping) else None + if isinstance(layer, Mapping): + parameters.update(layer) + if substitute and rw is not None and rw.environment: + parameters = _substitute(parameters, rw.environment) + return parameters + + +def typed_handler( + func: Callable[..., Any] | None = None, + *, + on_invalid: Literal["nack", "drop"] = "nack", + parameters_from: Sequence[str] | None = None, + substitute: bool = False, +) -> Any: + """Decorate a service callback to validate incoming messages. + + The decorated method keeps the ``(self, rw, header, message)`` call + signature expected by :func:`workflows.recipe.wrap_subscribe`. If the + message argument (the one after ``header``) is annotated with anything + other than ``Any``, the message is validated against that type and the + validated object passed on. If the method declares a ``parameters`` + argument, the parameters are assembled as described by parameters_from, + validated against its annotation, and passed in as a keyword argument. + + Can be used bare (``@typed_handler``) or with arguments + (``@typed_handler(parameters_from=("message", "step"))``). + + Args: + on_invalid: What to do with a message that fails validation. "nack" + rejects it without requeueing, so it goes to the dead letter queue + if one is configured. "drop" acknowledges and discards it. + parameters_from: Where the ``parameters`` argument is assembled from, + highest precedence first. Defaults to ``("step",)``, the recipe + step parameters only. See the module documentation for the + available sources and how they are merged. Only allowed if the + method declares a ``parameters`` argument. + substitute: Whether to replace $variables in the parameters with + values from the recipe environment. See the module documentation. + Only allowed if the method declares a ``parameters`` argument. + + Raises: + TypeError: If the method does not accept (self, rw, header, message), + or parameters_from or substitute is given for a method without a + ``parameters`` argument. + ValueError: If parameters_from is empty or names an unknown source. + """ + + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: + signature = inspect.signature(func) + positional = [ + p + for p in signature.parameters.values() + if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) + ] + if len(positional) < 4: + raise TypeError( + f"{func.__qualname__} must accept (self, rw, header, message) arguments" + ) + message_name = positional[3].name + takes_parameters = "parameters" in signature.parameters + if not takes_parameters: + for option, given in ( + ("parameters_from", parameters_from is not None), + ("substitute", substitute), + ): + if given: + raise TypeError( + f"{func.__qualname__} has no 'parameters' argument, so " + f"{option} has no effect" + ) + sources = ( + _DEFAULT_PARAMETERS_FROM if parameters_from is None else parameters_from + ) + # Check the sources now, so that mistakes fail at import time + _parse_sources(sources) + + # Resolve annotations on first use, so that with postponed evaluation + # models may be defined after the service class + adapters: dict[str, pydantic.TypeAdapter[Any] | None] = {} + + def get_adapters() -> dict[str, pydantic.TypeAdapter[Any] | None]: + if not adapters: + hints = typing.get_type_hints(func) + for name in (message_name, "parameters"): + hint = hints.get(name, inspect.Parameter.empty) + adapters[name] = ( + None if hint in _UNVALIDATED else pydantic.TypeAdapter(hint) + ) + return adapters + + @functools.wraps(func) + def wrapper( + self: Any, + rw: workflows.recipe.RecipeWrapper | None, + header: dict, + message: Any, + ) -> Any: + adapter = get_adapters() + log = getattr(self, "log", logger) + kwargs: dict[str, Any] = {} + try: + if takes_parameters: + parameters = merged_parameters( + rw, message, parameters_from=sources, substitute=substitute + ) + if param_adapter := adapter["parameters"]: + parameters = param_adapter.validate_python(parameters) + kwargs["parameters"] = parameters + if message_adapter := adapter[message_name]: + message = message_adapter.validate_python(message) + except pydantic.ValidationError as e: + log.error( + "Rejecting message failing validation in %s: %s", + func.__qualname__, + e, + ) + transport = rw.transport if rw is not None else self.transport + if on_invalid == "drop": + transport.ack(header) + else: + transport.nack(header, requeue=False) + return None + return func(self, rw, header, message, **kwargs) + + setattr(wrapper, workflows.recipe.TYPED_HANDLER_MARKER, True) + return wrapper + + if func is not None: + return decorator(func) + return decorator diff --git a/src/workflows/services/common_service.py b/src/workflows/services/common_service.py index acb3ad9c..7d3660ca 100644 --- a/src/workflows/services/common_service.py +++ b/src/workflows/services/common_service.py @@ -2,6 +2,7 @@ import contextlib import enum +import inspect import itertools import logging import multiprocessing.connection @@ -81,6 +82,36 @@ class Priority(enum.IntEnum): IDLE = 4 +def _check_recipe_callback(callback: Callable[..., Any]) -> None: + """Check that a callback can receive recipe messages. + + Recipe callbacks are called with (recipe_wrapper, header, message). A + callback that also requires a 'parameters' argument has to be decorated + with workflows.recipe.typed.typed_handler, which supplies it. Without + this check the mistake would only show up when the first message arrives. + """ + if getattr(callback, workflows.recipe.TYPED_HANDLER_MARKER, False): + return + try: + signature = inspect.signature(callback) + except (TypeError, ValueError): + # No signature available, eg. for some builtins + return + try: + signature.bind(None, None, None) + except TypeError: + name = getattr(callback, "__qualname__", repr(callback)) + hint = "" + if "parameters" in signature.parameters: + hint = ( + " To receive recipe parameters, decorate it with " + "workflows.recipe.typed.typed_handler." + ) + raise TypeError( + f"{name} can not be called with (rw, header, message) arguments.{hint}" + ) from None + + class CommonService: """ Base class for workflow services. A service is a piece of software that runs @@ -461,6 +492,7 @@ def subscribe( Returns: A unique subscription ID """ + _check_recipe_callback(callback) queue_name = self._resolve_queue(channel) if queue_name != channel: self.log.info("Subscribing to queue %s in place of %s", queue_name, channel) diff --git a/tests/recipe/test_typed.py b/tests/recipe/test_typed.py new file mode 100644 index 00000000..2ebca14e --- /dev/null +++ b/tests/recipe/test_typed.py @@ -0,0 +1,319 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Annotated, Any +from unittest import mock + +import pydantic +import pytest + +from workflows.recipe import Recipe, RecipeWrapper +from workflows.recipe.typed import NonEmpty, merged_parameters, typed_handler + + +class Payload(pydantic.BaseModel): + file: str + image: int = 1 + + +class Parameters(pydantic.BaseModel): + timeout: int = 60 + path: str + + +def make_rw(parameters: dict[str, Any], environment: dict[str, Any] | None = None): + recipe = Recipe( + { + 1: {"service": "test", "queue": "test", "parameters": parameters}, + "start": [(1, {})], + } + ) + return RecipeWrapper( + message={ + "recipe": recipe.recipe, + "recipe-pointer": 1, + "environment": environment or {}, + "payload": None, + }, + transport=mock.Mock(), + ) + + +class Service: + def __init__(self): + self.log = mock.Mock() + self.transport = mock.Mock() + self.calls: list[tuple] = [] + + @typed_handler + def untyped(self, rw, header, message: Any): + self.calls.append((rw, header, message)) + + @typed_handler + def typed(self, rw, header, message: Payload, parameters: Parameters): + self.calls.append((rw, header, message, parameters)) + + @typed_handler(substitute=True) + def substituting(self, rw, header, message: Payload, parameters: Parameters): + self.calls.append((rw, header, message, parameters)) + + @typed_handler(parameters_from=("message.parameters", "step")) + def overridable(self, rw, header, message: Payload, parameters: Parameters): + self.calls.append((rw, header, message, parameters)) + + @typed_handler(on_invalid="drop") + def dropping(self, rw, header, message: Payload): + self.calls.append((rw, header, message)) + + @typed_handler + def late(self, rw, header, message: DefinedLater): + self.calls.append((rw, header, message)) + + +class DefinedLater(pydantic.BaseModel): + value: int + + +def test_merged_parameters_defaults_to_step_only(): + rw = make_rw({"timeout": 10}) + assert merged_parameters(rw, {"timeout": 1, "parameters": {"timeout": 2}}) == { + "timeout": 10 + } + + +def test_merged_parameters_from_message_parameters(): + rw = make_rw({"timeout": 10, "path": "/"}) + assert merged_parameters( + rw, + {"timeout": 1, "parameters": {"timeout": 2}}, + parameters_from=("message.parameters", "step"), + ) == {"timeout": 2, "path": "/"} + + +def test_merged_parameters_from_whole_message(): + rw = make_rw({"timeout": 10, "path": "/"}) + assert merged_parameters( + rw, {"timeout": 1, "file": "a.h5"}, parameters_from=("message", "step") + ) == {"timeout": 1, "path": "/", "file": "a.h5"} + + +def test_merged_parameters_precedence_follows_order(): + rw = make_rw({"a": "step", "b": "step", "ispyb_parameters": {"a": "ispyb"}}) + message = {"a": "message", "b": "message"} + assert merged_parameters( + rw, message, parameters_from=("step.ispyb_parameters", "message", "step") + ) == {"a": "ispyb", "b": "message", "ispyb_parameters": {"a": "ispyb"}} + assert merged_parameters( + rw, message, parameters_from=("message", "step.ispyb_parameters", "step") + ) == {"a": "message", "b": "message", "ispyb_parameters": {"a": "ispyb"}} + + +@pytest.mark.parametrize("value", [0, "", False, None]) +def test_merged_parameters_falsy_values_take_precedence(value): + rw = make_rw({"a": "step"}) + assert merged_parameters(rw, {"a": value}, parameters_from=("message", "step")) == { + "a": value + } + + +def test_merged_parameters_merge_is_shallow(): + rw = make_rw({"cluster": {"memory": 4, "cpus": 2}}) + assert merged_parameters( + rw, {"cluster": {"cpus": 8}}, parameters_from=("message", "step") + ) == {"cluster": {"cpus": 8}} + + +@pytest.mark.parametrize( + "message", ["not a dict", [1, 2], None, {"parameters": "not a dict"}, {}] +) +def test_merged_parameters_skips_missing_sources(message): + rw = make_rw({"a": 1}) + assert merged_parameters( + rw, message, parameters_from=("message.parameters", "step") + ) == {"a": 1} + + +def test_merged_parameters_deeply_nested_source(): + assert merged_parameters( + None, + {"parameters": {"extra": {"a": 1}}}, + parameters_from=("message.parameters.extra",), + ) == {"a": 1} + + +def test_merged_parameters_without_recipe(): + message = {"parameters": {"a": 1}} + assert merged_parameters(None, message) == {} + assert merged_parameters( + None, message, parameters_from=("message.parameters", "step") + ) == {"a": 1} + + +def test_merged_parameters_does_not_substitute_by_default(): + rw = make_rw({"path": "/dls/$visit", "pid": "$$"}, environment={"visit": "cm1234"}) + assert merged_parameters(rw, {}) == {"path": "/dls/$visit", "pid": "$$"} + + +def test_merged_parameters_substitutes_environment(): + rw = make_rw( + {"path": "/dls/$visit", "nested": {"a": ["${visit}"]}, "other": "$unknown"}, + environment={"visit": "cm1234"}, + ) + assert merged_parameters( + rw, + {"from_message": "$visit"}, + parameters_from=("message", "step"), + substitute=True, + ) == { + "path": "/dls/cm1234", + "nested": {"a": ["cm1234"]}, + "other": "$unknown", + "from_message": "cm1234", + } + + +@pytest.mark.parametrize( + "sources,error", + [ + ("step", TypeError), + ((), ValueError), + (("recipe",), ValueError), + (("message.",), ValueError), + (("step..x",), ValueError), + ], +) +def test_merged_parameters_rejects_invalid_sources(sources, error): + with pytest.raises(error): + merged_parameters(None, {}, parameters_from=sources) + + +def test_handler_rejects_invalid_sources_at_decoration(): + with pytest.raises(ValueError): + + @typed_handler(parameters_from=("recipe",)) + def bad(self, rw, header, message, parameters): ... + + +def test_handler_rejects_sources_without_parameters_argument(): + with pytest.raises(TypeError, match="parameters_from"): + + @typed_handler(parameters_from=("message", "step")) + def bad(self, rw, header, message): ... + + +def test_handler_rejects_substitute_without_parameters_argument(): + with pytest.raises(TypeError, match="substitute"): + + @typed_handler(substitute=True) + def bad(self, rw, header, message): ... + + +def test_untyped_message_passed_through(): + service = Service() + rw = make_rw({}) + service.untyped(rw, {}, "anything") + assert service.calls == [(rw, {}, "anything")] + + +def test_valid_message_and_parameters_are_converted(): + service = Service() + rw = make_rw({"path": "/dls/$visit"}, environment={"visit": "cm1234"}) + service.typed(rw, {}, {"file": "a.h5", "parameters": {"timeout": "5"}}) + ((_, _, message, parameters),) = service.calls + assert message == Payload(file="a.h5") + assert parameters == Parameters(path="/dls/$visit", timeout=60) + + +def test_handler_substitutes_parameters(): + service = Service() + rw = make_rw({"path": "/dls/$visit"}, environment={"visit": "cm1234"}) + service.substituting(rw, {}, {"file": "a.h5"}) + ((_, _, _, parameters),) = service.calls + assert parameters == Parameters(path="/dls/cm1234") + + +def test_handler_parameters_from_message(): + service = Service() + rw = make_rw({"path": "/", "timeout": 10}) + service.overridable(rw, {}, {"file": "a.h5", "parameters": {"timeout": "5"}}) + ((_, _, _, parameters),) = service.calls + assert parameters == Parameters(path="/", timeout=5) + + +def test_invalid_message_is_nacked_without_requeue(): + service = Service() + rw = make_rw({"path": "/"}) + header = {"message-id": 1, "subscription": 1} + service.typed(rw, header, {"image": "x"}) + assert not service.calls + rw.transport.nack.assert_called_once_with(header, requeue=False) + service.log.error.assert_called_once() + + +def test_invalid_parameters_are_nacked(): + service = Service() + rw = make_rw({}) + service.typed(rw, {}, {"file": "a.h5"}) + assert not service.calls + rw.transport.nack.assert_called_once() + + +def test_invalid_message_can_be_dropped(): + service = Service() + rw = make_rw({}) + service.dropping(rw, {}, {}) + assert not service.calls + rw.transport.ack.assert_called_once_with({}) + rw.transport.nack.assert_not_called() + + +def test_non_recipe_message_rejected_on_service_transport(): + service = Service() + service.dropping(None, {}, {}) + service.transport.ack.assert_called_once_with({}) + + +def test_annotations_resolved_lazily(): + service = Service() + service.late(None, {}, {"value": "3"}) + assert service.calls == [(None, {}, DefinedLater(value=3))] + + +def test_handler_requires_message_argument(): + with pytest.raises(TypeError): + + @typed_handler + def bad(self, rw, header): ... + + +@pytest.mark.parametrize("value", ["", None, 3]) +def test_nonempty_path_rejects(value): + with pytest.raises(pydantic.ValidationError): + pydantic.TypeAdapter(Annotated[Path, NonEmpty]).validate_python(value) + + +def test_nonempty_path_converts(): + adapter = pydantic.TypeAdapter(Annotated[Path, NonEmpty]) + assert adapter.validate_python("a/b") == Path("a/b") + + +def test_nonempty_str(): + adapter = pydantic.TypeAdapter(Annotated[str, NonEmpty]) + assert adapter.validate_python("a") == "a" + with pytest.raises(pydantic.ValidationError): + adapter.validate_python("") + + +def test_typed_handler_accepted_by_subscribe(mocker): + from workflows.services.common_service import CommonService + + wrap_subscribe = mocker.patch("workflows.recipe.wrap_subscribe") + + class TypedService(CommonService): + @typed_handler + def receive(self, rw, header, message: Payload, parameters: Parameters): ... + + service = TypedService() + service.transport = mock.Mock() + service.subscribe("a", service.receive) + wrap_subscribe.assert_called_once() diff --git a/tests/services/test_common_service.py b/tests/services/test_common_service.py index 2f922c1d..bf6fded5 100644 --- a/tests/services/test_common_service.py +++ b/tests/services/test_common_service.py @@ -671,3 +671,40 @@ def test_subscribe_named_override_takes_precedence(wrap_subscribe): service.transport = mock.Mock() service.subscribe("a", mock.Mock()) assert wrap_subscribe.call_args.args[1] == "a.test" + + +class CallbackService(CommonService): + def three(self, rw, header, message): ... + + def varargs(self, *args, **kwargs): ... + + def with_parameters(self, rw, header, message, parameters): ... + + def with_extra(self, rw, header, message, extra): ... + + def with_optional(self, rw, header, message, parameters=None): ... + + +@pytest.mark.parametrize("name", ["three", "varargs", "with_optional"]) +def test_subscribe_accepts_recipe_callbacks(wrap_subscribe, name): + service = CallbackService() + service.transport = mock.Mock() + service.subscribe("a", getattr(service, name)) + wrap_subscribe.assert_called_once() + + +def test_subscribe_rejects_parameters_without_typed_handler(wrap_subscribe): + service = CallbackService() + service.transport = mock.Mock() + with pytest.raises(TypeError, match="typed_handler"): + service.subscribe("a", service.with_parameters) + wrap_subscribe.assert_not_called() + + +def test_subscribe_rejects_incompatible_callback(wrap_subscribe): + service = CallbackService() + service.transport = mock.Mock() + with pytest.raises(TypeError, match="can not be called") as e: + service.subscribe("a", service.with_extra) + assert "typed_handler" not in str(e.value) + wrap_subscribe.assert_not_called()