Refactor event stream tests with {client,server}IntegrationTests

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/protocols/eventstream/ClientEventStreamUnmarshallerGeneratorTest.kt

@@ -7,43 +7,60 @@ package software.amazon.smithy.rust.codegen.client.smithy.protocols.eventstream
 
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ArgumentsSource
-import software.amazon.smithy.codegen.core.Symbol
-import software.amazon.smithy.model.shapes.StructureShape
-import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenTarget
-import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
-import software.amazon.smithy.rust.codegen.core.smithy.generators.builderSymbol
-import software.amazon.smithy.rust.codegen.core.smithy.protocols.Protocol
-import software.amazon.smithy.rust.codegen.core.smithy.protocols.parse.EventStreamUnmarshallerGenerator
 import software.amazon.smithy.rust.codegen.core.testutil.EventStreamTestModels
-import software.amazon.smithy.rust.codegen.core.testutil.EventStreamTestTools
-import software.amazon.smithy.rust.codegen.core.testutil.EventStreamTestVariety
-import software.amazon.smithy.rust.codegen.core.testutil.TestEventStreamProject
+import software.amazon.smithy.rust.codegen.core.testutil.EventStreamUnmarshallTestCases.writeUnmarshallTestCases
+import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams
+import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
+import software.amazon.smithy.rust.codegen.core.testutil.unitTest
 
 class ClientEventStreamUnmarshallerGeneratorTest {
     @ParameterizedTest
     @ArgumentsSource(TestCasesProvider::class)
-    fun test(testCase: EventStreamTestModels.TestCase) {
-        EventStreamTestTools.runTestCase(
-            testCase,
-            object : ClientEventStreamBaseRequirements() {
-                override fun renderGenerator(
-                    codegenContext: ClientCodegenContext,
-                    project: TestEventStreamProject,
-                    protocol: Protocol,
-                ): RuntimeType {
-                    fun builderSymbol(shape: StructureShape): Symbol = shape.builderSymbol(codegenContext.symbolProvider)
-                    return EventStreamUnmarshallerGenerator(
-                        protocol,
-                        codegenContext,
-                        project.operationShape,
-                        project.streamShape,
-                        ::builderSymbol,
-                    ).render()
-                }
-            },
-            CodegenTarget.CLIENT,
-            EventStreamTestVariety.Unmarshall,
-        )
+    fun `test event stream unmarshaller generator`(testCase: EventStreamTestModels.TestCase) {
+        clientIntegrationTest(
+            testCase.model,
+            IntegrationTestParams(service = "test#TestService", addModuleToEventStreamAllowList = true),
+        ) { codegenContext, rustCrate ->
+            val crateName = codegenContext.moduleUseName()
+            val generator = "$crateName::event_stream_serde::TestStreamUnmarshaller"
+
+            rustCrate.integrationTest("unmarshall") {
+                writeUnmarshallTestCases(testCase, codegenTarget = CodegenTarget.CLIENT, generator, codegenContext)
+
+                unitTest(
+                    "unknown_message",
+                    """
+                    let message = msg("event", "NewUnmodeledMessageType", "application/octet-stream", b"hello, world!");
+                    let result = $generator::new().unmarshall(&message);
+                    assert!(result.is_ok(), "expected ok, got: {:?}", result);
+                    assert!(expect_event(result.unwrap()).is_unknown());
+                    """,
+                )
+
+                unitTest(
+                    "generic_error",
+                    """
+                    let message = msg(
+                        "exception",
+                        "UnmodeledError",
+                        "${testCase.responseContentType}",
+                        br#"${testCase.validUnmodeledError}"#
+                    );
+                    let result = $generator::new().unmarshall(&message);
+                    assert!(result.is_ok(), "expected ok, got: {:?}", result);
+                    match expect_error(result.unwrap()).kind {
+                        TestStreamErrorKind::Unhandled(err) => {
+                            let message = format!("{}", aws_smithy_types::error::display::DisplayErrorContext(&err));
+                            let expected = "message: \"unmodeled error\"";
+                            assert!(message.contains(expected), "Expected '{message}' to contain '{expected}'");
+                        }
+                        kind => panic!("expected generic error, but got {:?}", kind),
+                    }
+                    """,
+                )
+            }
+        }
     }
 }

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustType.kt

@@ -458,9 +458,11 @@ class Attribute(val inner: Writable) {
         val AllowDeprecated = Attribute(allow("deprecated"))
         val AllowIrrefutableLetPatterns = Attribute(allow("irrefutable_let_patterns"))
         val AllowUnreachableCode = Attribute(allow("unreachable_code"))
+        val AllowUnreachablePatterns = Attribute(allow("unreachable_patterns"))
         val AllowUnusedImports = Attribute(allow("unused_imports"))
         val AllowUnusedMut = Attribute(allow("unused_mut"))
         val AllowUnusedVariables = Attribute(allow("unused_variables"))
+        val AllowMissingDocs = Attribute(allow("missing_docs"))
         val CfgTest = Attribute(cfg("test"))
         val DenyMissingDocs = Attribute(deny("missing_docs"))
         val DocHidden = Attribute(doc("hidden"))

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt

@@ -116,7 +116,7 @@ private fun <T : AbstractCodeWriter<T>, U> T.withTemplate(
  * This enables conditionally wrapping a block in a prefix/suffix, e.g.
  *
  * ```
- * writer.withBlock("Some(", ")", conditional = symbol.isOptional()) {
+ * writer.conditionalBlock("Some(", ")", conditional = symbol.isOptional()) {
  *      write("symbolValue")
  * }
  * ```

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/EventStreamUnmarshallerGenerator.kt

@@ -21,8 +21,10 @@ import software.amazon.smithy.model.shapes.TimestampShape
 import software.amazon.smithy.model.shapes.UnionShape
 import software.amazon.smithy.model.traits.EventHeaderTrait
 import software.amazon.smithy.model.traits.EventPayloadTrait
+import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
 import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
+import software.amazon.smithy.rust.codegen.core.rustlang.Visibility
 import software.amazon.smithy.rust.codegen.core.rustlang.conditionalBlock
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
@@ -44,6 +46,16 @@ import software.amazon.smithy.rust.codegen.core.util.expectTrait
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
 import software.amazon.smithy.rust.codegen.core.util.toPascalCase
 
+fun RustModule.Companion.eventStreamSerdeModule(): RustModule.LeafModule =
+    RustModule.new(
+        "event_stream_serde",
+        visibility = Visibility.PUBLIC,
+        documentation = "TODO",
+        inline = false,
+        parent = RustModule.LibRs,
+        additionalAttributes = listOf(Attribute.AllowMissingDocs),
+    )
+
 class EventStreamUnmarshallerGenerator(
     private val protocol: Protocol,
     codegenContext: CodegenContext,
@@ -52,6 +64,7 @@ class EventStreamUnmarshallerGenerator(
     /** Function that maps a StructureShape into its builder symbol */
     private val builderSymbol: (StructureShape) -> Symbol,
 ) {
+    private val crateName = codegenContext.moduleUseName()
     private val model = codegenContext.model
     private val symbolProvider = codegenContext.symbolProvider
     private val codegenTarget = codegenContext.target
@@ -63,7 +76,7 @@ class EventStreamUnmarshallerGenerator(
         unionShape.eventStreamErrorSymbol(symbolProvider).toSymbol()
     }
     private val smithyEventStream = RuntimeType.smithyEventStream(runtimeConfig)
-    private val eventStreamSerdeModule = RustModule.private("event_stream_serde")
+    private val eventStreamSerdeModule = RustModule.eventStreamSerdeModule()
     private val codegenScope = arrayOf(
         "Blob" to RuntimeType.blob(runtimeConfig),
         "expect_fns" to smithyEventStream.resolve("smithy"),
@@ -87,15 +100,17 @@ class EventStreamUnmarshallerGenerator(
     }
 
     private fun RustWriter.renderUnmarshaller(unmarshallerType: RuntimeType, unionSymbol: Symbol) {
+        val unmarshallerTypeName = unmarshallerType.name
         rust(
             """
             ##[non_exhaustive]
             ##[derive(Debug)]
-            pub struct ${unmarshallerType.name};
+            pub struct $unmarshallerTypeName;
 
-            impl ${unmarshallerType.name} {
+            impl $unmarshallerTypeName {
+                /// Creates a new $unmarshallerTypeName
                 pub fn new() -> Self {
-                    ${unmarshallerType.name}
+                    $unmarshallerTypeName
                 }
             }
             """,
@@ -157,6 +172,7 @@ class EventStreamUnmarshallerGenerator(
                         "Output" to unionSymbol,
                         *codegenScope,
                     )
+
                     false -> rustTemplate(
                         "return Err(#{Error}::unmarshalling(format!(\"unrecognized :event-type: {}\", _unknown_variant)));",
                         *codegenScope,
@@ -182,6 +198,7 @@ class EventStreamUnmarshallerGenerator(
                     *codegenScope,
                 )
             }
+
             payloadOnly -> {
                 withBlock("let parsed = ", ";") {
                     renderParseProtocolPayload(unionMember)
@@ -192,6 +209,7 @@ class EventStreamUnmarshallerGenerator(
                     *codegenScope,
                 )
             }
+
             else -> {
                 rust("let mut builder = #T::default();", builderSymbol(unionStruct))
                 val payloadMember = unionStruct.members().firstOrNull { it.hasTrait<EventPayloadTrait>() }
@@ -268,6 +286,7 @@ class EventStreamUnmarshallerGenerator(
                     is BlobShape -> {
                         rustTemplate("#{Blob}::new(message.payload().as_ref())", *codegenScope)
                     }
+
                     is StringShape -> {
                         rustTemplate(
                             """
@@ -278,6 +297,7 @@ class EventStreamUnmarshallerGenerator(
                             *codegenScope,
                         )
                     }
+
                     is UnionShape, is StructureShape -> {
                         renderParseProtocolPayload(member)
                     }
@@ -315,6 +335,7 @@ class EventStreamUnmarshallerGenerator(
                     *codegenScope,
                 )
             }
+
             CodegenTarget.SERVER -> {}
         }
 
@@ -355,10 +376,15 @@ class EventStreamUnmarshallerGenerator(
                                 )
                             }
                         }
+
                         CodegenTarget.SERVER -> {
                             val target = model.expectShape(member.target, StructureShape::class.java)
                             val parser = protocol.structuredDataParser(operationShape).errorParser(target)
-                            val mut = if (parser != null) { " mut" } else { "" }
+                            val mut = if (parser != null) {
+                                " mut"
+                            } else {
+                                ""
+                            }
                             rust("let$mut builder = #T::default();", builderSymbol(target))
                             if (parser != null) {
                                 rustTemplate(
@@ -396,6 +422,7 @@ class EventStreamUnmarshallerGenerator(
             CodegenTarget.CLIENT -> {
                 rustTemplate("Ok(#{UnmarshalledMessage}::Error(#{OpError}::generic(generic)))", *codegenScope)
             }
+
             CodegenTarget.SERVER -> {
                 rustTemplate(
                     """
@@ -411,6 +438,6 @@ class EventStreamUnmarshallerGenerator(
 
     private fun UnionShape.eventStreamUnmarshallerType(): RuntimeType {
         val symbol = symbolProvider.toSymbol(this)
-        return RuntimeType("crate::event_stream_serde::${symbol.name.toPascalCase()}Unmarshaller")
+        return RuntimeType("$crateName::event_stream_serde::${symbol.name.toPascalCase()}Unmarshaller")
     }
 }

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/EventStreamErrorMarshallerGenerator.kt

@@ -26,6 +26,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.generators.error.eventStreamErrorSymbol
 import software.amazon.smithy.rust.codegen.core.smithy.generators.renderUnknownVariant
 import software.amazon.smithy.rust.codegen.core.smithy.generators.unknownVariantError
+import software.amazon.smithy.rust.codegen.core.smithy.protocols.parse.eventStreamSerdeModule
 import software.amazon.smithy.rust.codegen.core.smithy.rustType
 import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticEventStreamUnionTrait
 import software.amazon.smithy.rust.codegen.core.smithy.transformers.eventStreamErrors
@@ -50,7 +51,7 @@ class EventStreamErrorMarshallerGenerator(
     } else {
         unionShape.eventStreamErrorSymbol(symbolProvider).toSymbol()
     }
-    private val eventStreamSerdeModule = RustModule.private("event_stream_serde")
+    private val eventStreamSerdeModule = RustModule.eventStreamSerdeModule()
     private val errorsShape = unionShape.expectTrait<SyntheticEventStreamUnionTrait>()
     private val codegenScope = arrayOf(
         "MarshallMessage" to smithyEventStream.resolve("frame::MarshallMessage"),

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/EventStreamMarshallerGenerator.kt

@@ -38,6 +38,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.generators.UnionGenerator
 import software.amazon.smithy.rust.codegen.core.smithy.generators.renderUnknownVariant
 import software.amazon.smithy.rust.codegen.core.smithy.generators.unknownVariantError
 import software.amazon.smithy.rust.codegen.core.smithy.isOptional
+import software.amazon.smithy.rust.codegen.core.smithy.protocols.parse.eventStreamSerdeModule
 import software.amazon.smithy.rust.codegen.core.smithy.rustType
 import software.amazon.smithy.rust.codegen.core.util.dq
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
@@ -53,7 +54,7 @@ open class EventStreamMarshallerGenerator(
     private val payloadContentType: String,
 ) {
     private val smithyEventStream = RuntimeType.smithyEventStream(runtimeConfig)
-    private val eventStreamSerdeModule = RustModule.private("event_stream_serde")
+    private val eventStreamSerdeModule = RustModule.eventStreamSerdeModule()
     private val codegenScope = arrayOf(
         "MarshallMessage" to smithyEventStream.resolve("frame::MarshallMessage"),
         "Message" to smithyEventStream.resolve("frame::Message"),

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/testutil/EventStreamTestModels.kt

@@ -19,6 +19,7 @@ private fun fillInBaseModel(
 ): String = """
     namespace test
 
+    use smithy.framework#ValidationException
     use aws.protocols#$protocolName
 
     union TestUnion {
@@ -69,12 +70,28 @@ private fun fillInBaseModel(
         MessageWithNoHeaderPayloadTraits: MessageWithNoHeaderPayloadTraits,
         SomeError: SomeError,
     }
-    structure TestStreamInputOutput { @httpPayload @required value: TestStream }
+
+    structure TestStreamInput {
+        @httpLabel
+        @required
+        id: String,
+
+        @httpPayload
+        value: TestStream,
+    }
+
+    structure TestStreamOutput {
+        @httpPayload
+        value: TestStream
+    }
+
+    @http(method: "POST", uri: "/test/{id}")
     operation TestStreamOp {
-        input: TestStreamInputOutput,
-        output: TestStreamInputOutput,
-        errors: [SomeError],
+        input: TestStreamInput,
+        output: TestStreamOutput,
+        errors: [SomeError, ValidationException],
     }
+
     $extraServiceAnnotations
     @$protocolName
     service TestService { version: "123", operations: [TestStreamOp] }

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/testutil/EventStreamTestTools.kt

@@ -30,7 +30,6 @@ import software.amazon.smithy.rust.codegen.core.smithy.protocols.Protocol
 import software.amazon.smithy.rust.codegen.core.smithy.transformers.EventStreamNormalizer
 import software.amazon.smithy.rust.codegen.core.smithy.transformers.OperationNormalizer
 import software.amazon.smithy.rust.codegen.core.testutil.EventStreamMarshallTestCases.writeMarshallTestCases
-import software.amazon.smithy.rust.codegen.core.testutil.EventStreamUnmarshallTestCases.writeUnmarshallTestCases
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
 import software.amazon.smithy.rust.codegen.core.util.lookup
 import software.amazon.smithy.rust.codegen.core.util.outputShape
@@ -104,7 +103,7 @@ object EventStreamTestTools {
         test.project.lib {
             when (variety) {
                 EventStreamTestVariety.Marshall -> writeMarshallTestCases(testCase, generator)
-                EventStreamTestVariety.Unmarshall -> writeUnmarshallTestCases(testCase, codegenTarget, generator)
+                EventStreamTestVariety.Unmarshall -> Unit // writeUnmarshallTestCases(testCase, codegenTarget, generator)
             }
         }
         test.project.compileAndTest()