Fix assembler addTraits for some model sources

Author: kstich

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.

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()