Skip to content

Commit

Permalink
fix: case schemas with unions. Closes #700 (#723)
Browse files Browse the repository at this point in the history
  • Loading branch information
@marcosschroh
marcosschroh authored Aug 21, 2024
1 parent 6541274 commit 8ef31ff
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 0 deletions.
4 changes: 4 additions & 0 deletions dataclasses_avroschema/case.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,10 @@ def case_item(item: typing.Dict, case_type: str) -> typing.Dict:
# means that it is a complex type with a record
new_record = case_record(value, case_type=case_type)
new_field[key] = new_record
elif isinstance(value, list):
new_field[key] = [
case_record(element, case_type=case_type) if isinstance(element, dict) else element for element in value
]
else:
new_field[key] = value

Expand Down
80 changes: 80 additions & 0 deletions tests/case/test_case_records.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,58 @@
]


CASE_TO_DATA_WITH_UNION = [
(
case.CAMELCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "tigerAnimal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "name", "type": "string"}, {"name": "createdBy", "type": "string"}]}], "default": null}, {"name": "whaleFish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "name", "type": "string"}, {"name": "finsCount", "type": "long"}]}], "default": null}, {"name": "lionAnimal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.CAPITALCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "Tiger_animal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "Name", "type": "string"}, {"name": "Created_by", "type": "string"}]}], "default": null}, {"name": "Whale_fish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "Name", "type": "string"}, {"name": "Fins_count", "type": "long"}]}], "default": null}, {"name": "Lion_animal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.CONSTCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "TIGER_ANIMAL", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "NAME", "type": "string"}, {"name": "CREATED_BY", "type": "string"}]}], "default": null}, {"name": "WHALE_FISH", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "NAME", "type": "string"}, {"name": "FINS_COUNT", "type": "long"}]}], "default": null}, {"name": "LION_ANIMAL", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.LOWERCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "tiger_animal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "name", "type": "string"}, {"name": "created_by", "type": "string"}]}], "default": null}, {"name": "whale_fish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "name", "type": "string"}, {"name": "fins_count", "type": "long"}]}], "default": null}, {"name": "lion_animal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.PASCALCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "TigerAnimal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "Name", "type": "string"}, {"name": "CreatedBy", "type": "string"}]}], "default": null}, {"name": "WhaleFish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "Name", "type": "string"}, {"name": "FinsCount", "type": "long"}]}], "default": null}, {"name": "LionAnimal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.PATHCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "tiger/animal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "name", "type": "string"}, {"name": "created/by", "type": "string"}]}], "default": null}, {"name": "whale/fish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "name", "type": "string"}, {"name": "fins/count", "type": "long"}]}], "default": null}, {"name": "lion/animal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.SNAKECASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "tiger_animal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "name", "type": "string"}, {"name": "created_by", "type": "string"}]}], "default": null}, {"name": "whale_fish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "name", "type": "string"}, {"name": "fins_count", "type": "long"}]}], "default": null}, {"name": "lion_animal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.SPINALCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "tiger-animal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "name", "type": "string"}, {"name": "created-by", "type": "string"}]}], "default": null}, {"name": "whale-fish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "name", "type": "string"}, {"name": "fins-count", "type": "long"}]}], "default": null}, {"name": "lion-animal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.UPPERSPINALCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "TIGER-ANIMAL", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "NAME", "type": "string"}, {"name": "CREATED-BY", "type": "string"}]}], "default": null}, {"name": "WHALE-FISH", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "NAME", "type": "string"}, {"name": "FINS-COUNT", "type": "long"}]}], "default": null}, {"name": "LION-ANIMAL", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.TRIMCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "tiger_animal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "name", "type": "string"}, {"name": "created_by", "type": "string"}]}], "default": null}, {"name": "whale_fish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "name", "type": "string"}, {"name": "fins_count", "type": "long"}]}], "default": null}, {"name": "lion_animal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.UPPERCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "TIGER_ANIMAL", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "NAME", "type": "string"}, {"name": "CREATED_BY", "type": "string"}]}], "default": null}, {"name": "WHALE_FISH", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "NAME", "type": "string"}, {"name": "FINS_COUNT", "type": "long"}]}], "default": null}, {"name": "LION_ANIMAL", "type": ["null", "Cat"], "default": null}]}', # noqa
),
(
case.ALPHANUMCASE,
'{"type": "record", "name": "Animal", "fields": [{"name": "tigeranimal", "type": ["null", {"type": "record", "name": "Cat", "fields": [{"name": "name", "type": "string"}, {"name": "createdby", "type": "string"}]}], "default": null}, {"name": "whalefish", "type": ["null", {"type": "record", "name": "Fish", "fields": [{"name": "name", "type": "string"}, {"name": "finscount", "type": "long"}]}], "default": null}, {"name": "lionanimal", "type": ["null", "Cat"], "default": null}]}', # noqa
),
]


@pytest.mark.parametrize("case_type, schema", CASE_TO_DATA)
def test_case_record(case_type, schema):
class Event(AvroModel):
Expand Down Expand Up @@ -301,3 +353,31 @@ class Meta:

assert schema == UserAdvance.avro_schema(case_type=case_type)
assert json.loads(schema) == UserAdvance.avro_schema_to_python(case_type=case_type)


@pytest.mark.parametrize("case_type, schema", CASE_TO_DATA_WITH_UNION)
def test_case_with_union(case_type, schema):
class Cat(AvroModel):
name: str
created_by: str

class Meta:
schema_doc = False

class Fish(AvroModel):
name: str
fins_count: int

class Meta:
schema_doc = False

class Animal(AvroModel):
tiger_animal: typing.Optional[Cat] = None
whale_fish: typing.Optional[Fish] = None
lion_animal: typing.Optional[Cat] = None

class Meta:
schema_doc = False

assert schema == Animal.avro_schema(case_type=case_type)
assert json.loads(schema) == Animal.avro_schema_to_python(case_type=case_type)

0 comments on commit 8ef31ff

Please sign in to comment.