Fix creating @standardRegionalEndpoints nodes

smithy-lang · Mar 12, 2024 · 3d72963 · 3d72963
1 parent 45ac2fe
commit 3d72963
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/...nts/src/main/java/software/amazon/smithy/rulesengine/aws/traits/PartitionSpecialCase.java b/...nts/src/main/java/software/amazon/smithy/rulesengine/aws/traits/PartitionSpecialCase.java
@@ -6,6 +6,7 @@
 package software.amazon.smithy.rulesengine.aws.traits;
 
 import java.util.Objects;
+import java.util.Optional;
 import software.amazon.smithy.model.FromSourceLocation;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.node.Node;
@@ -64,10 +65,10 @@ public Boolean getFips() {
     @Override
     public Node toNode() {
         return Node.objectNodeBuilder()
-            .withMember(ENDPOINT, endpoint)
-            .withMember(DUAL_STACK, dualStack.toString())
-            .withMember(FIPS, fips.toString())
-            .build();
+                .withMember(ENDPOINT, endpoint)
+                .withOptionalMember(DUAL_STACK, Optional.ofNullable(dualStack).map(Node::from))
+                .withOptionalMember(FIPS, Optional.ofNullable(fips).map(Node::from))
+                .build();
     }
 
     @Override

diff --git a/...in/java/software/amazon/smithy/rulesengine/aws/traits/StandardRegionalEndpointsTrait.java b/...in/java/software/amazon/smithy/rulesengine/aws/traits/StandardRegionalEndpointsTrait.java
@@ -79,7 +79,7 @@ protected Node createNode() {
         return ObjectNode.objectNodeBuilder()
                 .sourceLocation(getSourceLocation())
                 .withMember(PARTITION_SPECIAL_CASES, partitionSpecialCasesNodeBuilder.build())
-                .withMember(REGION_SPECIAL_CASES, partitionSpecialCasesNodeBuilder.build())
+                .withMember(REGION_SPECIAL_CASES, regionSpecialCasesNodeBuilder.build())
                 .build();
     }
 

diff --git a/...ava/software/amazon/smithy/rulesengine/aws/traits/StandardRegionalEndpointsTraitTest.java b/...ava/software/amazon/smithy/rulesengine/aws/traits/StandardRegionalEndpointsTraitTest.java
@@ -6,6 +6,7 @@
 import java.util.List;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 class StandardRegionalEndpointsTraitTest {
@@ -25,12 +26,12 @@ public void loadsFromModel() {
 
         trait = getTraitFromService(model, "ns.foo#Service2");
 
-        assertEquals(trait.getRegionSpecialCases().size(), 0);
+        assertEquals(trait.getPartitionSpecialCases().size(), 0);
         assertEquals(trait.getRegionSpecialCases().size(), 0);
 
         trait = getTraitFromService(model, "ns.foo#Service3");
 
-        assertEquals(trait.getRegionSpecialCases().size(), 1);
+        assertEquals(trait.getPartitionSpecialCases().size(), 1);
         assertEquals(trait.getRegionSpecialCases().size(), 1);
         List<PartitionSpecialCase> partitionSpecialCases = trait.getPartitionSpecialCases().get("aws-us-gov");
 
@@ -46,6 +47,8 @@ public void loadsFromModel() {
 
         List<RegionSpecialCase> regionSpecialCases = trait.getRegionSpecialCases().get("us-east-1");
         assertEquals(regionSpecialCases.size(), 0);
+
+        Node.assertEquals(trait.toNode(), trait.toBuilder().build().toNode());
     }
 
     private StandardRegionalEndpointsTrait getTraitFromService(Model model, String service) {