Allow overriding AWS endpoints partitions

smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/language/functions/AwsPartition.java

@@ -5,6 +5,7 @@
 
 package software.amazon.smithy.rulesengine.aws.language.functions;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -24,6 +25,7 @@
 import software.amazon.smithy.rulesengine.language.syntax.expressions.functions.FunctionNode;
 import software.amazon.smithy.rulesengine.language.syntax.expressions.functions.LibraryFunction;
 import software.amazon.smithy.utils.MapUtils;
+import software.amazon.smithy.utils.SmithyInternalApi;
 import software.amazon.smithy.utils.SmithyUnstableApi;
 
 /**
@@ -41,24 +43,46 @@ public final class AwsPartition extends LibraryFunction {
     public static final Identifier INFERRED = Identifier.of("inferred");
 
     private static final Definition DEFINITION = new Definition();
-    private static final List<Partition> PARTITIONS;
+
+    // The following are mutable to allow for overriding the contents
+    // of the PARTITIONS list for test use cases. They MUST NOT have
+    // contents exposed directly, only through copies as is done in
+    // the `evaluate` method below.
+    private static final List<Partition> PARTITIONS = new ArrayList<>();
     private static final Map<String, Partition> REGION_MAP = new HashMap<>();
 
     static {
-        PARTITIONS = Partitions.fromNode(Node.parse(Partitions.class.getResourceAsStream("partitions.json")))
-                .getPartitions();
+        PARTITIONS.addAll(Partitions.fromNode(
+                        Node.parse(Partitions.class.getResourceAsStream("partitions.json")))
+                .getPartitions());
+        initializeRegionMap();
+    }
+
+    private AwsPartition(FunctionNode functionNode) {
+        super(DEFINITION, functionNode);
+    }
 
+    /**
+     * Overrides the partitions provided by default.
+     *
+     * @param partitions A list of partitions to set.
+     */
+    @SmithyInternalApi
+    public static void overridePartitions(Partitions partitions) {
+        PARTITIONS.clear();
+        PARTITIONS.addAll(partitions.getPartitions());
+        initializeRegionMap();
+    }
+
+    private static void initializeRegionMap() {
+        REGION_MAP.clear();
         for (Partition partition : PARTITIONS) {
             for (String region : partition.getRegions().keySet()) {
                 REGION_MAP.put(region, partition);
             }
         }
     }
 
-    private AwsPartition(FunctionNode functionNode) {
-        super(DEFINITION, functionNode);
-    }
-
     /**
      * Gets the {@link FunctionDefinition} implementation.
      *

smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/language/functions/AwsPartitionTest.java

@@ -0,0 +1,75 @@
+package software.amazon.smithy.rulesengine.aws.language.functions;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.not;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.rulesengine.aws.language.functions.partition.Partition;
+import software.amazon.smithy.rulesengine.aws.language.functions.partition.PartitionOutputs;
+import software.amazon.smithy.rulesengine.aws.language.functions.partition.Partitions;
+import software.amazon.smithy.rulesengine.language.evaluation.RuleEvaluator;
+import software.amazon.smithy.rulesengine.language.evaluation.value.RecordValue;
+import software.amazon.smithy.rulesengine.language.syntax.expressions.Expression;
+
+public class AwsPartitionTest {
+    @Test
+    public void eval() {
+        RecordValue result = evalWithRegion("us-west-2");
+
+        assertThat(result.get(AwsPartition.DNS_SUFFIX).expectStringValue().getValue(), not(equalTo("")));
+        assertThat(result.get(AwsPartition.DUAL_STACK_DNS_SUFFIX).expectStringValue().getValue(), not(equalTo("")));
+        assertThat(result.get(AwsPartition.SUPPORTS_FIPS).expectBooleanValue().getValue(), equalTo(true));
+        assertThat(result.get(AwsPartition.SUPPORTS_DUAL_STACK).expectBooleanValue().getValue(), equalTo(true));
+        assertThat(result.get(AwsPartition.IMPLICIT_GLOBAL_REGION).expectStringValue().getValue(), not(equalTo("")));
+    }
+
+    @Test
+    public void enumeratedRegionIsNotInferred() {
+        RecordValue result = evalWithRegion("us-west-1");
+
+        assertThat(result.get(AwsPartition.INFERRED).expectBooleanValue().getValue(), equalTo(false));
+    }
+
+    @Test
+    public void regionNotEnumeratedIsInferred() {
+        RecordValue result = evalWithRegion("us-west-3");
+
+        assertThat(result.get(AwsPartition.INFERRED).expectBooleanValue().getValue(), equalTo(true));
+    }
+
+    @Test
+    public void overridesPartitions() {
+        RecordValue result = evalWithRegion("us-west-1");
+        assertThat(result.get(AwsPartition.INFERRED).expectBooleanValue().getValue(), equalTo(false));
+
+        // Remove the enumerated regions, so the same region is now inferred.
+        AwsPartition.overridePartitions(Partitions.builder()
+                .addPartition(Partition.builder()
+                        .id("aws")
+                        .regionRegex("^(us|eu|ap|sa|ca|me|af)-\\w+-\\d+$")
+                        .outputs(PartitionOutputs.builder()
+                                .name("aws")
+                                .dnsSuffix("amazonaws.com")
+                                .dualStackDnsSuffix("api.aws")
+                                .supportsFips(true)
+                                .supportsDualStack(true)
+                                .implicitGlobalRegion("us-east-1")
+                                .build())
+                        .build())
+                .build());
+
+        result = evalWithRegion("us-west-1");
+        assertThat(result.get(AwsPartition.INFERRED).expectBooleanValue().getValue(), equalTo(true));
+
+        // Set the partitions back to what they were.
+        AwsPartition.overridePartitions(Partitions.fromNode(
+                        Node.parse(Partitions.class.getResourceAsStream("partitions.json"))));
+    }
+
+    private RecordValue evalWithRegion(String region) {
+        AwsPartition fn = AwsPartition.ofExpressions(Expression.of(region));
+        return fn.accept(new RuleEvaluator()).expectRecordValue();
+    }
+}