Fix toShapeId call in EndpointModifierIndex

smithy-lang · Nov 16, 2023 · bfb4c97 · bfb4c97
1 parent 79ad214
commit bfb4c97
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/...ts/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java b/...ts/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java
@@ -29,7 +29,7 @@ public EndpointModifierIndex(Model model) {
         for (ServiceShape serviceShape : model.getServiceShapes()) {
             Map<ShapeId, Trait> result = new TreeMap<>();
             for (Trait trait : serviceShape.getAllTraits().values()) {
-                Shape traitShape = model.getShape(trait.toShapeId()).get();
+                Shape traitShape = model.expectShape(trait.toShapeId());
                 if (traitShape.hasTrait(EndpointModifierTrait.ID)) {
                     result.put(trait.toShapeId(), trait);
                 }
@@ -49,6 +49,6 @@ public static EndpointModifierIndex of(Model model) {
      * @return Map of endpoint modifier trait ID to the trait
      */
     public Map<ShapeId, Trait> getEndpointModifierTraits(ToShapeId toShapeId) {
-        return endpointModifierTraits.get(toShapeId);
+        return endpointModifierTraits.get(toShapeId.toShapeId());
     }
 }
diff --git a/...rc/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java b/...rc/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java
@@ -24,6 +24,8 @@ public void loadsFromModel() {
         EndpointModifierIndex index = new EndpointModifierIndex(model);
 
         assertEquals(index.getEndpointModifierTraits(service1).size(), 1);
+        // Assert this works if you pass in the shape as well.
+        assertEquals(index.getEndpointModifierTraits(model.expectShape(service1)).size(), 1);
         assertInstanceOf(StandardRegionalEndpointsTrait.class,
                 index.getEndpointModifierTraits(service1).get(StandardRegionalEndpointsTrait.ID));