From 2ba24a925e69c4d9e456b0c338846aeea322c7ae Mon Sep 17 00:00:00 2001
From: kstich <kevin@kstich.com>
Date: Tue, 15 Aug 2023 15:17:35 -0700
Subject: [PATCH] Fix assembler addTraits for some model sources

---
 .../smithy/model/loader/ModelAssembler.java   | 11 ++---
 .../model/loader/ModelAssemblerTest.java      | 45 +++++++++++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/loader/ModelAssembler.java b/smithy-model/src/main/java/software/amazon/smithy/model/loader/ModelAssembler.java
index 03e9d6f417b..bff196c1d95 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/loader/ModelAssembler.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/loader/ModelAssembler.java
@@ -527,11 +527,6 @@ public ValidatedResult<Model> assemble() {
         // Register manually added shapes. Skip members because they are part of aggregate shapes.
         shapes.forEach(processor::putCreatedShape);
 
-        // Register manually added traits.
-        for (Pair<ShapeId, Trait> entry : pendingTraits) {
-            processor.accept(LoadOperation.ApplyTrait.from(entry.getKey(), entry.getValue()));
-        }
-
         // Register manually added Models.
         for (Model model : mergeModels) {
             // Add manually added metadata from the Model.
@@ -561,6 +556,12 @@ public ValidatedResult<Model> assemble() {
             }
         }
 
+        // Register manually added traits. Do this after loading any other sources of shapes
+        // so that traits can be applied to them.
+        for (Pair<ShapeId, Trait> entry : pendingTraits) {
+            processor.accept(LoadOperation.ApplyTrait.from(entry.getKey(), entry.getValue()));
+        }
+
         Model processedModel = processor.buildModel();
 
         // Do the 1.0 -> 2.0 transform before full-model validation.
diff --git a/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java b/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
index 9578cc80454..c50afdaaa93 100644
--- a/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
+++ b/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
@@ -130,6 +130,51 @@ public void addsExplicitTraits() {
         assertTrue(resultShape.getTrait(SuppressTrait.class).isPresent());
     }
 
+    @Test
+    public void addsExplicitTraitsToBuiltModel() {
+        StringShape shape = StringShape.builder().id("ns.foo#Bar").build();
+        SuppressTrait trait = SuppressTrait.builder().build();
+        ValidatedResult<Model> result = new ModelAssembler()
+                .addModel(Model.assembler().addShape(shape).assemble().unwrap())
+                .addTrait(shape.toShapeId(), trait)
+                .assemble();
+
+        assertThat(result.getValidationEvents(), empty());
+        Shape resultShape = result.unwrap().getShape(ShapeId.from("ns.foo#Bar")).get();
+        assertTrue(resultShape.findTrait("smithy.api#suppress").isPresent());
+        assertTrue(resultShape.getTrait(SuppressTrait.class).isPresent());
+    }
+
+    @Test
+    public void addsExplicitTraitsToUnparsedModel() {
+        String unparsed = "{\"smithy\": \"" + Model.MODEL_VERSION + "\", \"shapes\": { \"ns.foo#Bar\": { \"type\": \"string\"}}}";
+        SuppressTrait trait = SuppressTrait.builder().build();
+        ValidatedResult<Model> result = new ModelAssembler()
+                .addUnparsedModel(SourceLocation.NONE.getFilename(), unparsed)
+                .addTrait(ShapeId.from("ns.foo#Bar"), trait)
+                .assemble();
+
+        assertThat(result.getValidationEvents(), empty());
+        Shape resultShape = result.unwrap().getShape(ShapeId.from("ns.foo#Bar")).get();
+        assertTrue(resultShape.findTrait("smithy.api#suppress").isPresent());
+        assertTrue(resultShape.getTrait(SuppressTrait.class).isPresent());
+    }
+
+    @Test
+    public void addsExplicitTraitsToParsedDocumentNode() {
+        String unparsed = "{\"smithy\": \"" + Model.MODEL_VERSION + "\", \"shapes\": { \"ns.foo#Bar\": { \"type\": \"string\"}}}";
+        SuppressTrait trait = SuppressTrait.builder().build();
+        ValidatedResult<Model> result = new ModelAssembler()
+                .addDocumentNode(Node.parse(unparsed, SourceLocation.NONE.getFilename()))
+                .addTrait(ShapeId.from("ns.foo#Bar"), trait)
+                .assemble();
+
+        assertThat(result.getValidationEvents(), empty());
+        Shape resultShape = result.unwrap().getShape(ShapeId.from("ns.foo#Bar")).get();
+        assertTrue(resultShape.findTrait("smithy.api#suppress").isPresent());
+        assertTrue(resultShape.getTrait(SuppressTrait.class).isPresent());
+    }
+
     @Test
     public void addsExplicitDocumentNode_1_0_0() {
         ObjectNode node = Node.objectNode()