Fix incorrect builder properties assignment for enum with default value (#2752)

Author: joewyz

smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java

@@ -39,9 +39,12 @@
 import com.example.traits.numbers.LongTrait;
 import com.example.traits.numbers.ShortTrait;
 import com.example.traits.structures.BasicAnnotationTrait;
+import com.example.traits.structures.EnumA;
+import com.example.traits.structures.EnumB;
 import com.example.traits.structures.NestedA;
 import com.example.traits.structures.NestedB;
 import com.example.traits.structures.StructMemberWithTimestampFormatTrait;
+import com.example.traits.structures.StructWithEnumDefaultTrait;
 import com.example.traits.structures.StructWithIdrefMemberTrait;
 import com.example.traits.structures.StructWithListOfMapTrait;
 import com.example.traits.structures.StructWithUniqueItemsListTrait;
@@ -221,6 +224,12 @@ static Stream<Arguments> createTraitTests() {
                                 .memberEpochSeconds(Instant.ofEpochSecond((long) 1515531081.123))
                                 .build()
                                 .toNode()),
+                Arguments.of(StructWithEnumDefaultTrait.ID,
+                        StructWithEnumDefaultTrait.builder()
+                                .memberA(EnumA.TWO)
+                                .memberB(EnumB.FOUR)
+                                .build()
+                                .toNode()),
                 // Timestamps
                 Arguments.of(TimestampTrait.ID, Node.from("1985-04-12T23:20:50.52Z")),
                 Arguments.of(DateTimeTimestampTrait.ID, Node.from("1985-04-12T23:20:50.52Z")),

smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java

@@ -46,9 +46,12 @@
 import com.example.traits.numbers.LongTrait;
 import com.example.traits.numbers.ShortTrait;
 import com.example.traits.structures.BasicAnnotationTrait;
+import com.example.traits.structures.EnumA;
+import com.example.traits.structures.EnumB;
 import com.example.traits.structures.NestedA;
 import com.example.traits.structures.NestedB;
 import com.example.traits.structures.StructMemberWithTimestampFormatTrait;
+import com.example.traits.structures.StructWithEnumDefaultTrait;
 import com.example.traits.structures.StructWithIdrefMemberTrait;
 import com.example.traits.structures.StructWithListOfMapTrait;
 import com.example.traits.structures.StructWithUniqueItemsListTrait;
@@ -343,6 +346,13 @@ static Stream<Arguments> loadsModelTests() {
                                         DateTimeFormatter.RFC_1123_DATE_TIME.parse("Tue, 29 Apr 2014 18:30:38 GMT"))),
                                 "getMemberEpochSeconds",
                                 Optional.of(Instant.ofEpochSecond((long) 1515531081.123)))),
+                Arguments.of("structures/struct-with-enum-default-trait.smithy",
+                        StructWithEnumDefaultTrait.class,
+                        MapUtils.of(
+                                "getMemberA",
+                                EnumA.ONE,
+                                "getMemberB",
+                                EnumB.THREE)),
                 // Timestamps
                 Arguments.of("timestamps/struct-with-nested-timestamps.smithy",
                         StructWithNestedTimestampsTrait.class,

smithy-trait-codegen/src/it/resources/software/amazon/smithy/traitcodegen/test/structures/struct-with-enum-default-trait.smithy

@@ -0,0 +1,8 @@
+$version: "2.0"
+
+namespace test.smithy.traitcodegen
+
+use test.smithy.traitcodegen.structures#StructWithEnumDefault
+
+@StructWithEnumDefault
+structure myStruct {}

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java

@@ -21,7 +21,9 @@
 import software.amazon.smithy.model.shapes.ByteShape;
 import software.amazon.smithy.model.shapes.DocumentShape;
 import software.amazon.smithy.model.shapes.DoubleShape;
+import software.amazon.smithy.model.shapes.EnumShape;
 import software.amazon.smithy.model.shapes.FloatShape;
+import software.amazon.smithy.model.shapes.IntEnumShape;
 import software.amazon.smithy.model.shapes.IntegerShape;
 import software.amazon.smithy.model.shapes.ListShape;
 import software.amazon.smithy.model.shapes.LongShape;
@@ -534,5 +536,21 @@ public Void timestampShape(TimestampShape timestampShape) {
             writer.write("$T.parse($S)", Instant.class, defaultValue.expectStringNode().getValue());
             return null;
         }
+
+        @Override
+        public Void intEnumShape(IntEnumShape intEnumShape) {
+            writer.write("$T.from($L)",
+                    symbolProvider.toSymbol(member),
+                    defaultValue.expectNumberNode().getValue());
+            return null;
+        }
+
+        @Override
+        public Void enumShape(EnumShape enumShape) {
+            writer.write("$T.from($S)",
+                    symbolProvider.toSymbol(member),
+                    defaultValue.expectStringNode().getValue());
+            return null;
+        }
     }
 }

smithy-trait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java

@@ -26,7 +26,7 @@
 import software.amazon.smithy.model.node.ObjectNode;
 
 public class TraitCodegenPluginTest {
-    private static final int EXPECTED_NUMBER_OF_FILES = 69;
+    private static final int EXPECTED_NUMBER_OF_FILES = 72;
 
     private MockManifest manifest;
     private Model model;

smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest

@@ -44,6 +44,7 @@ structures/struct-with-listofmap-trait.smithy
 structures/struct-with-uniqueitems-list-trait.smithy
 structures/struct-with-idref-member-trait.smithy
 structures/struct-member-with-timestamp-format-trait.smithy
+structures/struct-with-enum-default-trait.smithy
 timestamps/date-time-format-timestamp-trait.smithy
 timestamps/epoch-seconds-format-timestamp-trait.smithy
 timestamps/http-date-format-timestamp-trait.smithy

smithy-trait-codegen/src/test/resources/META-INF/smithy/structures/struct-with-enum-default-trait.smithy

@@ -0,0 +1,19 @@
+$version: "2.0"
+
+namespace test.smithy.traitcodegen.structures
+
+@trait
+structure StructWithEnumDefault {
+    memberA: EnumA = "1"
+    memberB: EnumB = 3
+}
+
+enum EnumA {
+    ONE = "1"
+    TWO = "2"
+}
+
+intEnum EnumB {
+    THREE = 3
+    FOUR = 4
+}