Explode on unrecognized shape types in event stream test tools

smithy-lang · Dec 20, 2022 · c608a77 · c608a77
1 parent be77fa0
commit c608a77
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/...src/main/kotlin/software/amazon/smithy/rust/codegen/core/testutil/EventStreamTestTools.kt b/...src/main/kotlin/software/amazon/smithy/rust/codegen/core/testutil/EventStreamTestTools.kt
@@ -123,7 +123,7 @@ object EventStreamTestTools {
                 CodegenTarget.CLIENT -> CombinedErrorGenerator(model, symbolProvider, operationSymbol, errors).render(this)
                 CodegenTarget.SERVER -> ServerCombinedErrorGenerator(model, symbolProvider, operationSymbol, errors).render(this)
             }
-            for (shape in model.shapes().filter { shape -> shape.isStructureShape && shape.hasTrait<ErrorTrait>() }) {
+            for (shape in model.shapes().filter { shape -> shape is StructureShape && shape.hasTrait<ErrorTrait>() }) {
                 StructureGenerator(model, symbolProvider, this, shape as StructureShape).render(codegenTarget)
                 requirements.renderBuilderForShape(this, codegenContext, shape)
             }
@@ -155,10 +155,16 @@ object EventStreamTestTools {
         for (member in shape.members()) {
             val target = model.expectShape(member.target)
             if (target is StructureShape || target is UnionShape) {
-                if (target is StructureShape) {
-                    target.renderWithModelBuilder(model, symbolProvider, writer)
-                } else if (target is UnionShape) {
-                    UnionGenerator(model, symbolProvider, writer, target, renderUnknownVariant = mode.renderUnknownVariant()).render()
+                when (target) {
+                    is StructureShape -> target.renderWithModelBuilder(model, symbolProvider, writer)
+                    is UnionShape -> UnionGenerator(
+                        model,
+                        symbolProvider,
+                        writer,
+                        target,
+                        renderUnknownVariant = mode.renderUnknownVariant(),
+                    ).render()
+                    else -> TODO("EventStreamTestTools doesn't support rendering $target")
                 }
                 recursivelyGenerateModels(model, symbolProvider, target, writer, mode)
             }