EventStream Fixes (#1647)

aws · Mar 30, 2022 · f9446dc · f9446dc
1 parent b03e3da
commit f9446dc
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/...s-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsEventStreamUtils.java b/...s-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsEventStreamUtils.java
@@ -780,7 +780,7 @@ private static void generateEventStreamWriter(
         var writeCloser = getSymbol("WriteCloser", SmithyGoDependency.IO);
         var signerInterface = getModuleSymbol(settings, EVENT_STREAM_SIGNER_INTERFACE);
 
-        var messageSymbol = getEventStreamSymbol("Message", false);
+        var messageSymbol = getEventStreamSymbol("Message", true);
 
         if (withInitialMessages) {
             generateEventStreamWriterMessageWrapper(eventStream, service, symbolProvider, writer, eventUnionSymbol);
@@ -937,6 +937,9 @@ private static void generateEventStreamWriter(
                     writer.write("w.serializationBuffer.Reset()").write("")
                             .write("eventMessage := $T{}", messageSymbol).write("");
 
+                    var eventStreamSerializerName = getEventStreamSerializerName(eventStream, service,
+                            context.getProtocolName());
+
                     if (withInitialMessages) {
                         var initialRequestType = getWriterEventWrapperInitialRequestType(symbolProvider, eventStream,
                                 service);
@@ -956,14 +959,14 @@ private static void generateEventStreamWriter(
                                 default:
                                     return nil, $T("unknown event wrapper type: %v", event)
                                 }
-                                """, initialRequestType, messageEventType, errorf);
+                                """, initialRequestType, messageEventType, eventStreamSerializerName, errorf);
                     } else {
                         writer.write("""
                                         if err := $L(event, &eventMessage); err != nil {
                                             return nil, err
                                         }
                                         """,
-                                getEventStreamSerializerName(eventStream, service, context.getProtocolName()));
+                                eventStreamSerializerName);
                     }
 
                     writer.write("""