From 0a6c17aedad4e75e0985cfea4e0a2a629c28c7a9 Mon Sep 17 00:00:00 2001
From: Jonas Keeling <jonas.keeling@aiven.io>
Date: Mon, 6 May 2024 17:16:39 +0200
Subject: [PATCH] fix: flatten_unions for array and map union types

avro arrays and maps have no name property. Instead type property is used as name.
---
 karapace/serialization.py        |  2 ++
 tests/unit/test_serialization.py | 31 ++++++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/karapace/serialization.py b/karapace/serialization.py
index c199bad7a..3164283a5 100644
--- a/karapace/serialization.py
+++ b/karapace/serialization.py
@@ -422,6 +422,8 @@ def flatten_unions(schema: avro.schema.Schema, value: Any) -> Any:
         def get_name(obj) -> str:
             if isinstance(obj, avro.schema.PrimitiveSchema):
                 return obj.fullname
+            if isinstance(obj, (avro.schema.ArraySchema, avro.schema.MapSchema)):
+                return obj.type
             return obj.name
 
         f = next((s for s in schema.schemas if get_name(s) in value), None)
diff --git a/tests/unit/test_serialization.py b/tests/unit/test_serialization.py
index 2e09fec82..62567bc84 100644
--- a/tests/unit/test_serialization.py
+++ b/tests/unit/test_serialization.py
@@ -47,6 +47,18 @@
                     "name": "attr2",
                     "type": ["null", "string"],
                 },
+                {
+                    "name": "attrArray",
+                    "type": ["null", {"type": "array", "items": "string"}],
+                },
+                {
+                    "name": "attrMap",
+                    "type": ["null", {"type": "map", "values": "string"}],
+                },
+                {
+                    "name": "attrRecord",
+                    "type": ["null", {"type": "record", "name": "Record", "fields": [{"name": "attr1", "type": "string"}]}],
+                },
             ],
         }
     ),
@@ -128,13 +140,18 @@ async def test_happy_flow(default_config_path: Path):
     assert mock_registry_client.method_calls == [call.get_schema("top"), call.get_schema_for_id(1)]
 
 
-def test_flatten_unions_record() -> None:
-    record = {"attr1": {"string": "sample data"}, "attr2": None}
-    flatten_record = {"attr1": "sample data", "attr2": None}
-    assert flatten_unions(TYPED_AVRO_SCHEMA.schema, record) == flatten_record
-
-    record = {"attr1": None, "attr2": None}
-    assert flatten_unions(TYPED_AVRO_SCHEMA.schema, record) == record
+@pytest.mark.parametrize(
+    ["record", "flattened_record"],
+    [
+        [{"attr1": {"string": "sample data"}, "attr2": None}, {"attr1": "sample data", "attr2": None}],
+        [{"attr1": None, "attr2": None}, {"attr1": None, "attr2": None}],
+        [{"attrArray": {"array": ["item1", "item2"]}}, {"attrArray": ["item1", "item2"]}],
+        [{"attrMap": {"map": {"k1": "v1", "k2": "v2"}}}, {"attrMap": {"k1": "v1", "k2": "v2"}}],
+        [{"attrRecord": {"Record": {"attr1": "test"}}}, {"attrRecord": {"attr1": "test"}}],
+    ],
+)
+def test_flatten_unions_record(record, flattened_record) -> None:
+    assert flatten_unions(TYPED_AVRO_SCHEMA.schema, record) == flattened_record
 
 
 def test_flatten_unions_array() -> None: