1616package software .amazon .smithy .typescript .codegen .integration ;
1717
1818import java .util .List ;
19+ import java .util .Objects ;
1920import java .util .Optional ;
2021import java .util .Set ;
2122import java .util .TreeSet ;
4344import software .amazon .smithy .typescript .codegen .TypeScriptWriter ;
4445import software .amazon .smithy .typescript .codegen .integration .ProtocolGenerator .GenerationContext ;
4546import software .amazon .smithy .typescript .codegen .knowledge .SerdeElisionIndex ;
47+ import software .amazon .smithy .utils .Pair ;
4648import software .amazon .smithy .utils .SmithyUnstableApi ;
4749
4850/**
@@ -102,30 +104,38 @@ public void generateEventStreamSerializers(
102104 TopDownIndex topDownIndex = TopDownIndex .of (model );
103105 Set <OperationShape > operations = topDownIndex .getContainedOperations (service );
104106 TreeSet <UnionShape > eventUnionsToSerialize = new TreeSet <>();
105- TreeSet <StructureShape > eventShapesToMarshall = new TreeSet <>();
107+ TreeSet <Pair <String , StructureShape >> eventShapesToMarshall = new TreeSet <>(
108+ (a , b ) -> Objects .compare (a .getRight (), b .getRight (), StructureShape ::compareTo )
109+ );
110+
106111 for (OperationShape operation : operations ) {
107112 if (hasEventStreamInput (context , operation )) {
108113 UnionShape eventsUnion = getEventStreamInputShape (context , operation );
109114 eventUnionsToSerialize .add (eventsUnion );
110- Set <StructureShape > eventShapes = eventsUnion .members ().stream ()
111- .map (member -> model .expectShape (member .getTarget ()).asStructureShape ().get ())
112- .collect (Collectors .toSet ());
113- eventShapes .forEach (eventShapesToMarshall ::add );
115+ eventsUnion .members ()
116+ .forEach (member -> {
117+ eventShapesToMarshall .add (Pair .of (
118+ member .getMemberName (),
119+ model .expectShape (member .getTarget ()).asStructureShape ().get ()
120+ ));
121+ });
114122 }
115123 }
116124
117125 eventUnionsToSerialize .forEach (eventsUnion -> {
118126 generateEventStreamSerializer (context , eventsUnion );
119127 });
120128 SerdeElisionIndex serdeElisionIndex = SerdeElisionIndex .of (model );
121- eventShapesToMarshall .forEach (event -> {
129+ eventShapesToMarshall .forEach (memberNameAndEvent -> {
122130 generateEventMarshaller (
123131 context ,
124- event ,
132+ memberNameAndEvent .getLeft (),
133+ memberNameAndEvent .getRight (),
125134 documentContentType ,
126135 serializeInputEventDocumentPayload ,
127136 documentShapesToSerialize ,
128- serdeElisionIndex );
137+ serdeElisionIndex
138+ );
129139 });
130140 }
131141
@@ -155,6 +165,7 @@ public void generateEventStreamDeserializers(
155165 Set <OperationShape > operations = topDownIndex .getContainedOperations (service );
156166 TreeSet <UnionShape > eventUnionsToDeserialize = new TreeSet <>();
157167 TreeSet <StructureShape > eventShapesToUnmarshall = new TreeSet <>();
168+
158169 for (OperationShape operation : operations ) {
159170 if (hasEventStreamOutput (context , operation )) {
160171 UnionShape eventsUnion = getEventStreamOutputShape (context , operation );
@@ -237,6 +248,7 @@ private Symbol getSymbol(GenerationContext context, Shape shape) {
237248
238249 public void generateEventMarshaller (
239250 GenerationContext context ,
251+ String memberName ,
240252 StructureShape event ,
241253 String documentContentType ,
242254 Runnable serializeInputEventDocumentPayload ,
@@ -253,7 +265,7 @@ public void generateEventMarshaller(
253265 + "): __Message => {" , "}" , methodName , symbol , () -> {
254266 writer .openBlock ("const headers: __MessageHeaders = {" , "}" , () -> {
255267 //fix headers required by event stream
256- writer .write ("\" :event-type\" : { type: \" string\" , value: $S }," , symbol . getName () );
268+ writer .write ("\" :event-type\" : { type: \" string\" , value: $S }," , memberName );
257269 writer .write ("\" :message-type\" : { type: \" string\" , value: \" event\" }," );
258270 writeEventContentTypeHeader (context , event , documentContentType );
259271 });
0 commit comments