Add convenience method to get IntEnumShape from GenerateIntEnumDirective

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<C extends CodegenContext<S, ?, ?>, S> extends ShapeDirective<Shape, C, S> {
 
     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()));
     }
 }

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<TestContext, TestSettings>(null, null, shape);
+        });
+    }
+}