Fix nested dynamic wrapped structures and unions

Author: mtdowling

dynamic-schemas/build.gradle.kts

@@ -9,4 +9,6 @@ extra["moduleName"] = "software.amazon.smithy.java.dynamicschemas"
 
 dependencies {
     api(project(":core"))
+
+    testImplementation(project(":codecs:json-codec"))
 }

dynamic-schemas/src/main/java/software/amazon/smithy/java/dynamicschemas/SchemaInterceptingSerializer.java

@@ -13,8 +13,10 @@
 import software.amazon.smithy.java.core.schema.Schema;
 import software.amazon.smithy.java.core.schema.SerializableStruct;
 import software.amazon.smithy.java.core.serde.MapSerializer;
+import software.amazon.smithy.java.core.serde.SerializationException;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.java.core.serde.document.Document;
+import software.amazon.smithy.java.core.serde.document.DocumentParser;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ShapeType;
 
@@ -45,23 +47,38 @@ public <T> void writeList(Schema schema, T listState, int size, BiConsumer<T, Sh
 
     @Override
     public <T> void writeMap(Schema schema, T mapState, int size, BiConsumer<T, MapSerializer> consumer) {
-        var delegateKeySchema = delegateSchema.mapKeyMember();
-        var delegateValueSchema = delegateSchema.mapValueMember();
-        delegateSerializer.writeMap(delegateSchema, mapState, size, (s, ms) -> {
-            consumer.accept(s, new MapSerializer() {
-                @Override
-                public <V> void writeEntry(
-                        Schema keySchema,
-                        String key,
-                        V state,
-                        BiConsumer<V, ShapeSerializer> valueSerializer
-                ) {
-                    ms.writeEntry(delegateKeySchema, key, state, (s, ms) -> {
-                        valueSerializer.accept(s, new SchemaInterceptingSerializer(service, delegateValueSchema, ms));
+        switch (delegateSchema.type()) {
+            case DOCUMENT, MAP -> {
+                var delegateKeySchema = delegateSchema.mapKeyMember();
+                var delegateValueSchema = delegateSchema.mapValueMember();
+                delegateSerializer.writeMap(delegateSchema, mapState, size, (s, ms) -> {
+                    consumer.accept(s, new MapSerializer() {
+                        @Override
+                        public <V> void writeEntry(
+                                Schema keySchema,
+                                String key,
+                                V state,
+                                BiConsumer<V, ShapeSerializer> valueSerializer
+                        ) {
+                            ms.writeEntry(delegateKeySchema, key, state, (s, ms) -> {
+                                valueSerializer.accept(s,
+                                        new SchemaInterceptingSerializer(service, delegateValueSchema, ms));
+                            });
+                        }
                     });
-                }
-            });
-        });
+                });
+            }
+            case STRUCTURE, UNION -> {
+                var parser = new DocumentParser();
+                parser.writeMap(schema, mapState, size, consumer);
+                var struct = parser.getResult();
+                var wrappedStruct = new WrappedDocument(delegateSchema, Document.of(struct), service);
+                delegateSerializer.writeStruct(delegateSchema, wrappedStruct);
+            }
+            default -> {
+                throw new SerializationException("Expected a map, structure, or union, but found " + delegateSchema);
+            }
+        }
     }
 
     @Override

dynamic-schemas/src/test/java/software/amazon/smithy/java/dynamicschemas/SchemaGuidedDocumentBuilderTest.java

@@ -7,8 +7,6 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.not;
-import static org.hamcrest.Matchers.nullValue;
 
 import java.math.BigDecimal;
 import java.math.BigInteger;
@@ -24,15 +22,13 @@
 import org.junit.jupiter.params.provider.MethodSource;
 import software.amazon.smithy.java.core.schema.PreludeSchemas;
 import software.amazon.smithy.java.core.schema.Schema;
-import software.amazon.smithy.java.core.schema.SerializableStruct;
 import software.amazon.smithy.java.core.serde.InterceptingSerializer;
-import software.amazon.smithy.java.core.serde.ShapeDeserializer;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.java.core.serde.SpecificShapeDeserializer;
 import software.amazon.smithy.java.core.serde.SpecificShapeSerializer;
 import software.amazon.smithy.java.core.serde.document.Document;
 import software.amazon.smithy.java.core.serde.document.DocumentDeserializer;
-import software.amazon.smithy.java.core.serde.document.DocumentUtils;
+import software.amazon.smithy.java.json.JsonCodec;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ShapeType;
@@ -96,6 +92,16 @@ enum MyEnum {
                             foo: String
                             baz: SimpleStruct
                         }
+
+                        map StructMap {
+                            key: String
+                            value: Foo
+                        }
+
+                        structure Foo {
+                            @jsonName("B")
+                            b: String
+                        }
                         """)
                 .assemble()
                 .unwrap();
@@ -265,6 +271,26 @@ public void throwsWhenNoValueSetForUnion() {
         Assertions.assertThrows(IllegalArgumentException.class, builder::build);
     }
 
+    @Test
+    public void worksWithMapOfStructure() {
+        var source = Document.ofObject(Map.of("a", Document.ofObject(Map.of("b", "str"))));
+
+        var converter = new SchemaConverter(model);
+        var schema = converter.getSchema(model.expectShape(ShapeId.from("smithy.example#StructMap")));
+
+        var builder = SchemaConverter.createDocumentBuilder(schema, ShapeId.from("smithy.example#StructMap"));
+        source.deserializeInto(builder);
+
+        var result = builder.build();
+
+        assertThat(result.asObject(), equalTo(source.asObject()));
+        assertThat(result.schema(), equalTo(schema));
+
+        var codec = JsonCodec.builder().useJsonName(true).build();
+
+        assertThat(codec.serializeToString(result), equalTo("{\"a\":{\"B\":\"str\"}}"));
+    }
+
     static List<Arguments> deserializesShapesProvider() {
         return List.of(
                 Arguments.of("MyDocument", Document.ofObject(Map.of("a", "b"))),