From 706e8626f31222bfbaac5b77a12d636ffc477d12 Mon Sep 17 00:00:00 2001
From: David Piggott <david@piggott.me.uk>
Date: Tue, 20 Feb 2024 17:20:54 +0000
Subject: [PATCH] Change target type to Document

---
 .../META-INF/smithy/unknownfieldretention.smithy       |  2 +-
 .../core/test/resources/META-INF/smithy/traits.smithy  |  7 +------
 .../alloy/UnknownFieldRetentionTraitProviderSpec.scala | 10 ++--------
 3 files changed, 4 insertions(+), 15 deletions(-)

diff --git a/modules/core/resources/META-INF/smithy/unknownfieldretention.smithy b/modules/core/resources/META-INF/smithy/unknownfieldretention.smithy
index 89da186..0939658 100644
--- a/modules/core/resources/META-INF/smithy/unknownfieldretention.smithy
+++ b/modules/core/resources/META-INF/smithy/unknownfieldretention.smithy
@@ -4,6 +4,6 @@ namespace alloy
 
 /// Retain unknown fields of a containing structure in a map.
 @trait(
-    selector: "structure > member :test(> map :test(> member > document))"
+    selector: "structure > member :test(> document)"
 )
 structure unknownFieldRetention {}
diff --git a/modules/core/test/resources/META-INF/smithy/traits.smithy b/modules/core/test/resources/META-INF/smithy/traits.smithy
index 07ea3f8..c4e17e8 100644
--- a/modules/core/test/resources/META-INF/smithy/traits.smithy
+++ b/modules/core/test/resources/META-INF/smithy/traits.smithy
@@ -211,10 +211,5 @@ structure TestUnknownFieldRetention {
    foo: String
    bar: String
    @unknownFieldRetention
-   bazes: RetainedUnknownFields
-}
-
-map RetainedUnknownFields {
-   key: String
-   value: Document
+   bazes: Document
 }
diff --git a/modules/core/test/src/alloy/UnknownFieldRetentionTraitProviderSpec.scala b/modules/core/test/src/alloy/UnknownFieldRetentionTraitProviderSpec.scala
index fbd0ed7..f460b70 100644
--- a/modules/core/test/src/alloy/UnknownFieldRetentionTraitProviderSpec.scala
+++ b/modules/core/test/src/alloy/UnknownFieldRetentionTraitProviderSpec.scala
@@ -15,7 +15,6 @@
 
 package alloy
 
-import software.amazon.smithy.model.shapes.ListShape
 import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.shapes.DocumentShape
@@ -30,11 +29,6 @@ final class UnknownFieldRetentionTraitProviderSpec extends munit.FunSuite {
       .builder()
       .id(ShapeId.fromParts("test", "MyDocument"))
       .build()
-    val mapShape = ListShape
-      .builder()
-      .id(ShapeId.fromParts("test", "MyMap"))
-      .member(documentShape.getId)
-      .build()
     val structId = ShapeId.fromParts("test", "MyStruct")
     val targetId = structId.withMember("myMap")
     val structShape = StructureShape
@@ -44,7 +38,7 @@ final class UnknownFieldRetentionTraitProviderSpec extends munit.FunSuite {
         MemberShape
           .builder()
           .id(targetId)
-          .target(mapShape.getId)
+          .target(documentShape.getId)
           .addTrait(new UnknownFieldRetentionTrait)
           .build()
       )
@@ -52,7 +46,7 @@ final class UnknownFieldRetentionTraitProviderSpec extends munit.FunSuite {
 
     val model =
       Model.assembler.disableValidation
-        .addShapes(structShape, mapShape, documentShape)
+        .addShapes(structShape, documentShape)
         .assemble()
         .unwrap()