diff --git a/tests/materializer/unions/test_mixed_discrimination.py b/tests/materializer/unions/test_callable_string_mixed_discrimination.py similarity index 94% rename from tests/materializer/unions/test_mixed_discrimination.py rename to tests/materializer/unions/test_callable_string_mixed_discrimination.py index 60f34a7..d4bbf97 100644 --- a/tests/materializer/unions/test_mixed_discrimination.py +++ b/tests/materializer/unions/test_callable_string_mixed_discrimination.py @@ -1,4 +1,4 @@ -"""Tests for mixed string/callable discriminated unions. +"""Tests for mixed callable/string discriminated unions with a simple model union member. Covers: - Callable discriminated union where one tag point to simple model @@ -116,7 +116,7 @@ class NestedGroupedThing(BaseModel): @pytest.mark.parametrize("param", params) -def test_mixed_discriminated_union(param): +def test_callable_string_mixed_discrimination(param): materializer = ModelMaterializer(**param.kwargs) assert list(materializer.generate_bindings()) == param.expected.bindings assert [ @@ -124,7 +124,7 @@ def test_mixed_discriminated_union(param): ] == param.expected.model_dump -def test_mixed_discriminated_union_model_types(): +def test_callable_string_mixed_discrimination_model_types(): cat, dog, car = ModelMaterializer(model=Model, data=data).generate_models() assert isinstance(cat.thing, Cat) assert isinstance(dog.thing, Dog) diff --git a/tests/materializer/unions/test_string_callable_mixed_discrimination.py b/tests/materializer/unions/test_string_callable_mixed_discrimination.py new file mode 100644 index 0000000..a284353 --- /dev/null +++ b/tests/materializer/unions/test_string_callable_mixed_discrimination.py @@ -0,0 +1,188 @@ +"""Tests for mixed string/callable discriminated unions with a simple model union member. + +Covers: +- String discriminated union where one member is a simple model + and another points to a nested callable discriminated union. + This is the inverse test to test_callable_string_mixed_discrimination.py. +- Discrimnated union model aggregation (top-level + nested union aggregation) +""" + +from typing import Annotated, Any, Literal + +import pytest +from graphty import ConfigDict, ModelMaterializer +from pydantic import BaseModel, Discriminator, Field, Tag +from tests.materializer.param import Expected, Parameter + + +class Cat(BaseModel): + kind: Literal["cat"] + name: str + + +class RoadBike(BaseModel): + kind: Literal["bike"] + bike_type: Literal["road"] + + +class MountainBike(BaseModel): + kind: Literal["bike"] + bike_type: Literal["mountain"] + + +def get_bike_type(v: Any) -> str | None: + bike_type = ( + v.get("bike_type") if isinstance(v, dict) else getattr(v, "bike_type", None) + ) + if bike_type == "road": + return "road" + if bike_type == "mountain": + return "mountain" + + +Bike = Annotated[ + Annotated[RoadBike, Tag("road")] | Annotated[MountainBike, Tag("mountain")], + Discriminator(get_bike_type), +] + +Thing = Annotated[Bike | Cat, Field(discriminator="kind")] + + +class Model(BaseModel): + thing: Thing + + +class GroupedThing(BaseModel): + model_config = ConfigDict(group_by="kind") + + kind: str + nested: list[Thing] + + +class NestedGroupedThing(BaseModel): + model_config = ConfigDict(group_by="kind") + + kind: str + nested: list[GroupedThing] + + +data = [ + {"kind": "cat", "name": "Muzi", "bike_type": None}, + {"kind": "bike", "name": None, "bike_type": "road"}, + {"kind": "bike", "name": None, "bike_type": "mountain"}, +] + + +params = [ + Parameter( + kwargs={"model": Model, "data": data}, + expected=Expected( + bindings=[ + {"thing": {"kind": "cat", "bike_type": None, "name": "Muzi"}}, + {"thing": {"kind": "bike", "bike_type": "road", "name": None}}, + {"thing": {"kind": "bike", "bike_type": "mountain", "name": None}}, + ], + model_dump=[ + {"thing": {"kind": "cat", "name": "Muzi"}}, + {"thing": {"kind": "bike", "bike_type": "road"}}, + {"thing": {"kind": "bike", "bike_type": "mountain"}}, + ], + ), + ), + Parameter( + kwargs={"model": GroupedThing, "data": data}, + expected=Expected( + bindings=[ + { + "kind": "cat", + "nested": [{"kind": "cat", "bike_type": None, "name": "Muzi"}], + }, + { + "kind": "bike", + "nested": [ + {"kind": "bike", "bike_type": "road", "name": None}, + {"kind": "bike", "bike_type": "mountain", "name": None}, + ], + }, + ], + model_dump=[ + {"kind": "cat", "nested": [{"kind": "cat", "name": "Muzi"}]}, + { + "kind": "bike", + "nested": [ + {"kind": "bike", "bike_type": "road"}, + {"kind": "bike", "bike_type": "mountain"}, + ], + }, + ], + ), + ), + Parameter( + kwargs={"model": NestedGroupedThing, "data": data}, + expected=Expected( + bindings=[ + { + "kind": "cat", + "nested": [ + { + "kind": "cat", + "nested": [ + {"kind": "cat", "bike_type": None, "name": "Muzi"} + ], + } + ], + }, + { + "kind": "bike", + "nested": [ + { + "kind": "bike", + "nested": [ + {"kind": "bike", "bike_type": "road", "name": None}, + {"kind": "bike", "bike_type": "mountain", "name": None}, + ], + } + ], + }, + ], + model_dump=[ + { + "kind": "cat", + "nested": [ + {"kind": "cat", "nested": [{"kind": "cat", "name": "Muzi"}]} + ], + }, + { + "kind": "bike", + "nested": [ + { + "kind": "bike", + "nested": [ + {"kind": "bike", "bike_type": "road"}, + {"kind": "bike", "bike_type": "mountain"}, + ], + } + ], + }, + ], + ), + ), +] + + +@pytest.mark.parametrize("param", params) +def test_string_callable_mixed_discrimination(param): + materializer = ModelMaterializer(**param.kwargs) + assert list(materializer.generate_bindings()) == param.expected.bindings + assert [ + m.model_dump() for m in materializer.generate_models() + ] == param.expected.model_dump + + +def test_string_callable_mixed_discrimination_model_types(): + cat, road_bike, mountain_bike = ModelMaterializer( + model=Model, data=data + ).generate_models() + assert isinstance(cat.thing, Cat) + assert isinstance(road_bike.thing, RoadBike) + assert isinstance(mountain_bike.thing, MountainBike)