Repository: pydantic/pydantic-ai
File: tests/test_temporal.py (Python)

Code before:

    pass


class CustomPayloadConverter(DefaultPayloadConverter):
    """A custom payload converter that does not inherit from PydanticPayloadConverter."""

    pass


class MockPayloadCodec(PayloadCodec):
    """A mock payload codec for testing (simulates encryption codec)."""

    async def encode(
        self, payloads: Sequence[temporalio.api.common.v1.Payload]
    ) -> list[temporalio.api.common.v1.Payload]:  # pragma: no cover
        return list(payloads)

    async def decode(
        self, payloads: Sequence[temporalio.api.common.v1.Payload]
    ) -> list[temporalio.api.common.v1.Payload]:  # pragma: no cover
        return list(payloads)


def test_pydantic_ai_json_payload_converter_preserves_document_url_media_type() -> None:
    """Test that `PydanticAIJSONPayloadConverter` preserves `DocumentUrl.media_type` during round-trip.

    This is a regression test for https://github.com/pydantic/pydantic-ai/issues/3949
    """

    converter = PydanticAIPayloadConverter()
    json_converter = converter.converters[b'json/plain']

    original = DocumentUrl(url='https://example.com/doc/12345', media_type='application/pdf')

    payload = json_converter.to_payload(original)
    assert payload is not None

    restored = json_converter.from_payload(payload, DocumentUrl)  # pyright: ignore[reportUnknownMemberType]
    assert restored == snapshot(
        DocumentUrl(url='https://example.com/doc/12345', _media_type='application/pdf', _identifier='eb8998')
    )


def test_pydantic_ai_plugin_no_converter_returns_pydantic_ai_data_converter() -> None:
    """When no converter is provided, PydanticAIPlugin uses pydantic_ai_data_converter."""
    plugin = PydanticAIPlugin()
    config: dict[str, Any] = {}
    result = plugin.configure_client(config)  # type: ignore[arg-type]
    assert result['data_converter'] is pydantic_ai_data_converter


Review comment (suggestion):
I'd rather just have the higher level test, so we can potentially swap out the implementation while continuing to verify that using a DocumentUrl (with media type) works as expected.

Code after:

class MockPayloadCodec(PayloadCodec):
    """A mock payload codec for testing (simulates encryption codec)."""

    async def encode(
        self, payloads: Sequence[temporalio.api.common.v1.Payload]
    ) -> list[temporalio.api.common.v1.Payload]:  # pragma: no cover
        return list(payloads)

    async def decode(
        self, payloads: Sequence[temporalio.api.common.v1.Payload]
    ) -> list[temporalio.api.common.v1.Payload]:  # pragma: no cover
        return list(payloads)


def test_pydantic_ai_plugin_no_converter_returns_pydantic_data_converter() -> None:
    """When no converter is provided, PydanticAIPlugin uses the standard pydantic_data_converter."""
    plugin = PydanticAIPlugin()
    # Create a minimal config without data_converter
    config: dict[str, Any] = {}
    result = plugin.configure_client(config)  # type: ignore[arg-type]
    assert result['data_converter'] is pydantic_data_converter


def test_pydantic_ai_plugin_with_pydantic_payload_converter_unchanged() -> None:
    """When converter already uses PydanticPayloadConverter, return it unchanged."""
    plugin = PydanticAIPlugin()
    converter = DataConverter(payload_converter_class=PydanticPayloadConverter)
    config: dict[str, Any] = {'data_converter': converter}
    result = plugin.configure_client(config)  # type: ignore[arg-type]
    assert result['data_converter'] is converter


def test_pydantic_ai_plugin_with_custom_pydantic_subclass_unchanged() -> None:
    """When converter uses a subclass of PydanticPayloadConverter, return it unchanged (no warning)."""
    plugin = PydanticAIPlugin()
    converter = DataConverter(payload_converter_class=CustomPydanticPayloadConverter)
    config: dict[str, Any] = {'data_converter': converter}
    result = plugin.configure_client(config)  # type: ignore[arg-type]
    assert result['data_converter'] is converter
    assert result['data_converter'].payload_converter_class is CustomPydanticPayloadConverter


def test_pydantic_ai_plugin_with_default_payload_converter_replaced() -> None:
    """When converter uses DefaultPayloadConverter, replace payload_converter_class without warning."""
    plugin = PydanticAIPlugin()
    converter = DataConverter(payload_converter_class=DefaultPayloadConverter)
    config: dict[str, Any] = {'data_converter': converter}
    result = plugin.configure_client(config)  # type: ignore[arg-type]
    assert result['data_converter'] is not converter
    assert result['data_converter'].payload_converter_class is PydanticPayloadConverter