from typing import Annotated, Literal

from pydantic import BaseModel, Field, WrapValidator


class ObjectPayload(BaseModel):
    name: str


class AddMutation(BaseModel):
    type: Literal["add"] = Field("add")
    # passes if you replace WrapValidator with AfterValidator or remove it
    payload: Annotated[ObjectPayload, WrapValidator(lambda v, nxt: nxt(v))] = Field()


class UpdateMutation(BaseModel):
    type: Literal["update"] = Field("update")
    payload: ObjectPayload = Field()


class ObjectUpdate(BaseModel):
    # passes if you use Field(union_mode="left_to_right")
    mutation: AddMutation | UpdateMutation = Field()


def test_union():
    obj = ObjectUpdate.model_validate(
        {
            "mutation": {
                "type": "add",
                "payload": {"name": "John Doe"},
            }
        }
    )
    assert obj.mutation.type == "add"
    assert obj.mutation.payload.name == "John Doe"

    obj = ObjectUpdate.model_validate(
        {
            "mutation": {
                "type": "update",
                "payload": {"name": "John Doe"},
            }
        }
    )
    assert obj.mutation.type == "update"
    assert obj.mutation.payload.name == "John Doe"

    obj = ObjectUpdate.model_validate(
        {
            "mutation": {
                "payload": {"name": "John Doe"},
            }
        }
    )
    assert obj.mutation.type == "add"
    assert obj.mutation.payload.name == "John Doe"