diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java index 1ab87b7f7ba..dc3518af691 100644 --- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java +++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java @@ -16,6 +16,8 @@ package software.amazon.smithy.codegen.core.directed; import software.amazon.smithy.codegen.core.CodegenContext; +import software.amazon.smithy.model.node.ExpectationNotMetException; +import software.amazon.smithy.model.shapes.IntEnumShape; import software.amazon.smithy.model.shapes.ServiceShape; import software.amazon.smithy.model.shapes.Shape; @@ -29,6 +31,19 @@ public final class GenerateIntEnumDirective, S> extends ShapeDirective { GenerateIntEnumDirective(C context, ServiceShape service, Shape shape) { - super(context, service, shape); + super(context, service, validateShape(shape)); + } + + private static Shape validateShape(Shape shape) { + if (shape.isIntEnumShape()) { + return shape; + } + throw new IllegalArgumentException("GenerateIntEnum requires an IntEnum shape"); + } + + + public IntEnumShape expectIntEnumShape() { + return shape().asIntEnumShape().orElseThrow(() -> new ExpectationNotMetException( + "Expected an IntEnum shape, but found " + shape(), shape())); } } diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirectiveTest.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirectiveTest.java new file mode 100644 index 00000000000..2d582fbabf0 --- /dev/null +++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirectiveTest.java @@ -0,0 +1,16 @@ +package software.amazon.smithy.codegen.core.directed; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.model.shapes.BooleanShape; + +public class GenerateIntEnumDirectiveTest { + @Test + public void validatesShapeType() { + BooleanShape shape = BooleanShape.builder().id("smithy.example#Foo").build(); + + Assertions.assertThrows(IllegalArgumentException.class, () -> { + new GenerateIntEnumDirective(null, null, shape); + }); + } +}