From bfb4c97653b14874227be0d9d83214d233b721e3 Mon Sep 17 00:00:00 2001 From: kstich Date: Thu, 16 Nov 2023 10:52:48 -0800 Subject: [PATCH] Fix toShapeId call in EndpointModifierIndex --- .../smithy/rulesengine/aws/traits/EndpointModifierIndex.java | 4 ++-- .../rulesengine/aws/traits/EndpointModifierIndexTest.java | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java b/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java index 9f65c647c3f..ff8ccbc59f1 100644 --- a/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java +++ b/smithy-aws-endpoints/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 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 getEndpointModifierTraits(ToShapeId toShapeId) { - return endpointModifierTraits.get(toShapeId); + return endpointModifierTraits.get(toShapeId.toShapeId()); } } diff --git a/smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java b/smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java index 53e79c0c14e..fa11edd67aa 100644 --- a/smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java +++ b/smithy-aws-endpoints/src/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));