Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4398260636
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -140,11 +217,215 @@ class UnresolvedForwardReference:
value: Missing # type: ignore[name-defined] # noqa: F821
+class ItemPair(NamedTuple):
+ item: Item
+
+
class Point(NamedTuple):
x: int
y: str
+class TestIsModelDumpRequired:
+ @pytest.mark.parametrize(
+ ("output", "output_type", "expected"),
+ [
+ (A(x=1), [A, int], True),
+ (A(x=1), (A, int), True),
+ (A(x=1), ToolOutput(A), True),
+ (A(x=1), A, False),
+ (A(x=1), A | int, False),
+ (A(x=1), list[A], False),
+ (SubA(x=1), A, True),
+ (5, [A, int], False),
+ ],
+ ids=[
+ "list-of-types",
+ "tuple-of-types",
+ "output-marker",
+ "single-model",
+ "union",
+ "generic-list",
+ "subclass-instance",
+ "not-a-model",
+ ],
+ )
+ @requires_walker
+ def test_is_true_when_the_worker_registers_no_class_for_the_output(self,
output, output_type, expected):
+ result = is_model_dump_required(output, output_type)
+
+ assert result is expected
+
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def test_is_false_without_the_sdk_walker(self):
+ assert is_model_dump_required(A(x=1), [A, int]) is False
+
+
+class TestIsModelDumpRequiredForContainerOutputs:
+ @pytest.mark.parametrize(
+ ("output", "output_type"),
+ [
+ ([A(x=1)], list[A]),
+ ({"k": A(x=1)}, dict[str, A]),
+ (None, A | None),
+ ("text", [list[Item], str]),
+ ],
+ ids=["registered-list", "registered-dict", "none-for-optional",
"no-model-in-output"],
+ )
+ @requires_walker
+ def
test_is_false_for_a_registered_shape_or_an_output_without_a_model(self, output,
output_type):
+ assert is_model_dump_required(output, output_type) is False
+
+ @pytest.mark.parametrize(
+ ("output", "output_type"),
+ [
+ ([Item(n=1)], [list[Item], str]),
+ ({"k": Item(n=1)}, [dict[str, Item], str]),
+ ((Item(n=1),), [tuple[Item, ...], str]),
+ (ItemPair(item=Item(n=1)), [ItemPair, str]),
+ ],
+ ids=["list", "dict", "tuple", "named-tuple"],
+ )
+ @requires_walker
+ def
test_is_true_for_models_inside_a_container_the_worker_registers_no_class_for(
+ self, output, output_type
+ ):
+ assert is_model_dump_required(output, output_type) is True
+
+
+class TestDumpOutputForXcom:
+ @requires_walker
+ def
test_dumped_model_with_http_url_enum_path_serializes_through_serde(self):
+ page = Page(url="https://example.com/a";, kind=Color.RED,
path=Path("/tmp/a"))
+
+ dumped = dump_output_for_xcom(page, [Page, str],
serialize_output=False)
+
+ assert deserialize(serialize(dumped)) == {
+ "url": "https://example.com/a";,
+ "kind": "red",
+ "path": "/tmp/a",
+ }
+
+ @requires_walker
+ def
test_dumps_models_inside_a_list_when_output_type_is_a_list_of_types(self):
+ dumped = dump_output_for_xcom([Item(n=1)], [list[Item], str],
serialize_output=False)
+
+ assert dumped == [{"n": 1}]
+ assert deserialize(serialize(dumped)) == [{"n": 1}]
+
+ @requires_walker
+ def test_keeps_a_registered_container_as_model_instances(self):
+ output = {"k": A(x=1)}
+
+ dumped = dump_output_for_xcom(output, dict[str, A],
serialize_output=False)
+
+ assert dumped is output
+ assert isinstance(dumped["k"], A)
+
+ @pytest.mark.parametrize(
+ "output",
+ [datetime(2026, 1, 1, tzinfo=UTC), 5, "s", None, {"a": 1}],
+ ids=["datetime", "int", "str", "none", "dict"],
+ )
+ def test_non_model_outputs_are_returned_unchanged(self, output):
+ dumped = dump_output_for_xcom(output, [A, int], serialize_output=True)
+
+ assert dumped is output
+
+ @requires_walker
+ def
test_dumps_the_whole_output_when_only_some_of_its_models_are_unregistered(self):
+ output = [A(x=1), Item(n=1)]
+
+ assert is_model_dump_required(output, A) is True
+ assert dump_output_for_xcom(output, A, serialize_output=False) ==
[{"x": 1}, {"n": 1}]
+
+ @pytest.mark.parametrize("model_cls", [SerializationAlias,
SerializedByAlias])
+ @requires_walker
+ def test_dumps_field_names_not_aliases_like_dump_output_to_json(self,
model_cls):
+ model = model_cls(user_name="kax")
+
+ dumped = dump_output_for_xcom(model, [model_cls, str],
serialize_output=False)
+
+ assert dumped == {"user_name": "kax"}
+ assert dumped == json.loads(dump_output_to_json(model))
+
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def test_keeps_a_container_without_the_sdk_walker(self):
+ output = [A(x=1)]
+
+ assert dump_output_for_xcom(output, list[A], serialize_output=True) is
output
+ assert dump_output_for_xcom(output, list[A], serialize_output=False)
is output
+
+ def test_serialize_output_keeps_python_field_types(self):
+ dumped = dump_output_for_xcom(STAMPED, Stamped, serialize_output=True)
+
+ assert type(dumped) is dict
+ assert dumped == {"d": STAMPED.d, "at": STAMPED.at, "amount":
STAMPED.amount, "ident": STAMPED.ident}
+ assert [type(value) for value in dumped.values()] == [date, datetime,
Decimal, UUID]
+
+ @pytest.mark.parametrize(
+ ("output", "expected"),
+ [
+ ([STAMPED], [STAMPED_JSON]),
+ ({"k": STAMPED}, {"k": STAMPED_JSON}),
+ ((STAMPED,), [STAMPED_JSON]),
+ ({STAMPED}, [STAMPED_JSON]),
+ ],
+ ids=["list", "dict", "tuple", "set"],
+ )
+ @requires_walker
+ def
test_serialize_output_dumps_an_unregistered_container_to_json_values(self,
output, expected):
+ dumped = dump_output_for_xcom(output, [Stamped, str],
serialize_output=True)
+
+ assert dumped == expected
+
+ @requires_walker
+ def test_dumps_a_model_used_as_a_dict_key(self):
+ output = {Key(n=1): 1, "k": 2}
+
+ dumped = dump_output_for_xcom(output, [Key, str],
serialize_output=True)
+
+ assert dumped == {"n=1": 1, "k": 2}
+ assert json.dumps(dumped)
+
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def
test_serialize_output_dumps_a_top_level_model_without_the_sdk_walker(self):
+ dumped = dump_output_for_xcom(STAMPED, Stamped, serialize_output=True)
+
+ assert type(dumped) is dict
+ assert type(dumped["d"]) is date
+
+ @pytest.mark.parametrize(
+ ("output", "output_type", "expected"),
+ [
+ ([LINK], list[Link], [LINK_SERIALIZED]),
+ ({"k": LINK}, dict[str, Link], {"k": LINK_SERIALIZED}),
+ ],
+ ids=["list", "dict"],
+ )
+ @requires_walker
+ def
test_serialize_output_keeps_registered_models_in_a_container_that_serde_pushes(
+ self, output, output_type, expected
+ ):
+ from airflow.sdk.serde import serialize
Review Comment:
Updated it as `MagicMock(spec=RunContext)` and removed the local import.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
