Fixing bug with enums as eventHeaders (#4337)

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here -->
Fixing a bug that was discovered with smithy `enum` and `intEnum` shapes
that are marked as `@eventHeader`. This came up in an internal
investigation around a failed preview build.

## Description
<!--- Describe your changes in detail -->
Added handling of `EnumShape` in `EventStreamUnmarshallerGenerator` and
`EventStreamMarshallerGenerator`.

## Testing
<!--- Please describe in detail how you tested your changes -->
<!--- Include details of your testing environment, and the tests you ran
to -->
<!--- see how your change affects other areas of the code, etc. -->
Added new `enum` and `intEnum` shapes to the evenstream test model. They
are only included in client tests since the server throws a validation
error on enums in event streams (see
smithy-lang/smithy#1388 for more info).

Also tested against the model from the failed preview build and a `cargo
test --all-features` was successful.


----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

Author: landonxjames

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

@@ -155,7 +155,11 @@ class ClientEventStreamMarshallerGeneratorTest {
                                             .is_some());
                                         assert_eq!(
                                             msg.payload(),
-                                            &bytes::Bytes::from_static(${testCase.generateRustPayloadInitializer(rpcEventStreamTestCase.expectedInInitialRequest)})
+                                            &bytes::Bytes::from_static(${
+                                            testCase.generateRustPayloadInitializer(
+                                                rpcEventStreamTestCase.expectedInInitialRequest,
+                                            )
+                                        })
                                         );
                                         """,
                                     )
@@ -178,7 +182,7 @@ class ClientEventStreamMarshallerGeneratorTest {
 
 class TestCasesProvider : ArgumentsProvider {
     override fun provideArguments(context: ExtensionContext?): Stream<out Arguments> =
-        EventStreamTestModels.TEST_CASES.map { Arguments.of(it) }.stream()
+        EventStreamTestModels.TEST_CASES.map { Arguments.of(it.withEnumMembers()) }.stream()
 }
 
 /**

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

@@ -9,6 +9,7 @@ import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.model.shapes.BlobShape
 import software.amazon.smithy.model.shapes.BooleanShape
 import software.amazon.smithy.model.shapes.ByteShape
+import software.amazon.smithy.model.shapes.EnumShape
 import software.amazon.smithy.model.shapes.IntegerShape
 import software.amazon.smithy.model.shapes.LongShape
 import software.amazon.smithy.model.shapes.MemberShape
@@ -261,6 +262,7 @@ class EventStreamUnmarshallerGenerator(
                     is IntegerShape -> rustTemplate("#{expect_fns}::expect_int32(header)?", *codegenScope)
                     is LongShape -> rustTemplate("#{expect_fns}::expect_int64(header)?", *codegenScope)
                     is BlobShape -> rustTemplate("#{expect_fns}::expect_byte_array(header)?", *codegenScope)
+                    is EnumShape -> rustTemplate("#{expect_fns}::expect_string(header)?.as_str().into()", *codegenScope)
                     is StringShape -> rustTemplate("#{expect_fns}::expect_string(header)?", *codegenScope)
                     is TimestampShape -> rustTemplate("#{expect_fns}::expect_timestamp(header)?", *codegenScope)
                     else -> throw IllegalStateException("unsupported event stream header shape type: $target")
@@ -383,7 +385,8 @@ class EventStreamUnmarshallerGenerator(
                                             "builder", target,
                                             mapErr = {
                                                 rustTemplate(
-                                                    """|err|#{Error}::unmarshalling(format!("{}", err))""", *codegenScope,
+                                                    """|err|#{Error}::unmarshalling(format!("{}", err))""",
+                                                    *codegenScope,
                                                 )
                                             },
                                         ),

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

@@ -10,6 +10,7 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.BlobShape
 import software.amazon.smithy.model.shapes.BooleanShape
 import software.amazon.smithy.model.shapes.ByteShape
+import software.amazon.smithy.model.shapes.EnumShape
 import software.amazon.smithy.model.shapes.IntegerShape
 import software.amazon.smithy.model.shapes.LongShape
 import software.amazon.smithy.model.shapes.MemberShape
@@ -243,6 +244,7 @@ open class EventStreamMarshallerGenerator(
             is IntegerShape -> "Int32($inputName)"
             is LongShape -> "Int64($inputName)"
             is BlobShape -> "ByteArray($inputName.into_inner().into())"
+            is EnumShape -> "String($inputName.to_string().into())"
             is StringShape -> "String($inputName.into())"
             is TimestampShape -> "Timestamp($inputName)"
             else -> throw IllegalStateException("unsupported event stream header shape type: $target")

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

@@ -21,8 +21,11 @@ private fun fillInBaseModel(
     namespacedProtocolName: String,
     extraServiceAnnotations: String = "",
     nonEventStreamMembers: String = "",
+    extraEventHeaderMembers: String = "",
+    extraShapes: String = "",
 ): String =
     """
+    ${"\$version: \"2\""}
     namespace test
 
     use smithy.framework#ValidationException
@@ -42,6 +45,8 @@ private fun fillInBaseModel(
         Message: String,
     }
 
+    $extraShapes
+
     structure MessageWithBlob { @eventPayload data: Blob }
     structure MessageWithString { @eventPayload data: String }
     structure MessageWithStruct { @eventPayload someStruct: TestStruct }
@@ -55,6 +60,7 @@ private fun fillInBaseModel(
         @eventHeader short: Short,
         @eventHeader string: String,
         @eventHeader timestamp: Timestamp,
+        $extraEventHeaderMembers
     }
     structure MessageWithHeaderAndPayload {
         @eventHeader header: String,
@@ -129,6 +135,26 @@ object EventStreamTestModels {
 
         fun withNonEventStreamMembers(nonEventStreamMembers: String): TestCase =
             this.copy(model = fillInBaseModel(this.protocolShapeId, "", nonEventStreamMembers).asSmithyModel())
+
+        // Server doesn't support enum members in event streams, so this util allows Clients to test with those shapes
+        fun withEnumMembers(): TestCase =
+            this.copy(
+                model =
+                    fillInBaseModel(
+                        this.protocolShapeId,
+                        extraEventHeaderMembers = "@eventHeader enum: TheEnum,\n@eventHeader intEnum: FaceCard,",
+                        extraShapes = """    enum TheEnum {
+                            FOO
+                            BAR
+                            }
+
+                            intEnum FaceCard {
+                            JACK = 1
+                            QUEEN = 2
+                            KING = 3
+                                            }""",
+                    ).asSmithyModel(),
+            )
     }
 
     private fun base64Encode(input: ByteArray): String {