fix: locate event stream member more carefully

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/EventStreamGenerator.java

@@ -39,6 +39,7 @@
 import software.amazon.smithy.model.traits.ErrorTrait;
 import software.amazon.smithy.model.traits.EventHeaderTrait;
 import software.amazon.smithy.model.traits.EventPayloadTrait;
+import software.amazon.smithy.model.traits.HttpPayloadTrait;
 import software.amazon.smithy.model.traits.StreamingTrait;
 import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
@@ -82,6 +83,26 @@ public static UnionShape getEventStreamOutputShape(GenerationContext context, Op
         return eventStreamInfo.getEventStreamTarget().asUnionShape().get();
     }
 
+    public static MemberShape getEventStreamMember(GenerationContext context, StructureShape struct) {
+        List<MemberShape> eventStreamMembers = struct.members()
+            .stream()
+            .filter(shape -> {
+                Shape target = context.getModel().expectShape(shape.getTarget());
+                boolean targetStreaming = target.hasTrait(StreamingTrait.class);
+                boolean targetUnion = target.isUnionShape();
+                boolean memberStreaming = shape.hasTrait(StreamingTrait.class);
+                boolean memberPayload = shape.hasTrait(HttpPayloadTrait.class);
+                return memberPayload && targetUnion && (targetStreaming || memberStreaming);
+            }).toList();
+
+        if (eventStreamMembers.isEmpty()) {
+            throw new CodegenException("No event stream member found in " + struct.getId().toString());
+        } else if (eventStreamMembers.size() > 1) {
+            throw new CodegenException("More than one event stream member in " + struct.getId().toString());
+        }
+        return eventStreamMembers.get(0);
+    }
+
     /**
      * Generate eventstream serializers, and related serializers for events.
      * @param context Code generation context instance.

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/HttpRpcProtocolGenerator.java

@@ -18,7 +18,6 @@
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.logging.Logger;
-import java.util.stream.Collectors;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.codegen.core.SymbolReference;
@@ -359,12 +358,13 @@ protected boolean writeRequestBody(GenerationContext context, OperationShape ope
             // Write the default `body` property.
             writer.write("let body: any;");
             if (EventStreamGenerator.hasEventStreamInput(context, operation)) {
-                // There must only one eventstream member in request structure.
-                MemberShape member = inputShape.members().stream().collect(Collectors.toList()).get(0);
-                Shape target = context.getModel().expectShape(member.getTarget());
+                MemberShape eventStreamMember = EventStreamGenerator.getEventStreamMember(
+                    context, inputShape
+                );
+                Shape target = context.getModel().expectShape(eventStreamMember.getTarget());
                 Symbol targetSymbol = context.getSymbolProvider().toSymbol(target);
                 String serFunctionName = ProtocolGenerator.getSerFunctionShortName(targetSymbol);
-                String memberName = member.getMemberName();
+                String memberName = eventStreamMember.getMemberName();
                 writer.write("body = $L(input.$L, context);", serFunctionName, memberName);
             } else {
                 // Track input shapes so their serializers may be generated.
@@ -540,12 +540,16 @@ protected void readResponseBody(GenerationContext context, OperationShape operat
                     // If there's an output present, we know it's a structure.
                     StructureShape outputShape = context.getModel().expectShape(outputId).asStructureShape().get();
                     if (EventStreamGenerator.hasEventStreamOutput(context, operation)) {
-                        // There must only one eventstream member in response structure.
-                        MemberShape member = outputShape.members().stream().collect(Collectors.toList()).get(0);
-                        Shape target = context.getModel().expectShape(member.getTarget());
+                        MemberShape eventStreamMember = EventStreamGenerator.getEventStreamMember(
+                            context, outputShape
+                        );
+                        Shape target = context.getModel().expectShape(eventStreamMember.getTarget());
                         Symbol targetSymbol = context.getSymbolProvider().toSymbol(target);
-                        writer.write("const contents = { $L: $L(output.body, context) };", member.getMemberName(),
-                                ProtocolGenerator.getDeserFunctionShortName(targetSymbol));
+                        writer.write(
+                            "const contents = { $L: $L(output.body, context) };",
+                            eventStreamMember.getMemberName(),
+                            ProtocolGenerator.getDeserFunctionShortName(targetSymbol)
+                        );
                     } else {
                         // We only need to load the body and prepare a contents object if there is a response.
                         writer.write("const data: any = await parseBody(output.body, context)");

smithy-typescript-codegen/src/test/java/software/amazon/smithy/typescript/codegen/integration/EventStreamGeneratorTest.java

@@ -0,0 +1,115 @@
+package software.amazon.smithy.typescript.codegen.integration;
+
+import java.util.List;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import software.amazon.smithy.codegen.core.CodegenException;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.shapes.StructureShape;
+import software.amazon.smithy.model.shapes.UnionShape;
+import software.amazon.smithy.model.traits.HttpPayloadTrait;
+import software.amazon.smithy.model.traits.StreamingTrait;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+class EventStreamGeneratorTest {
+    @Test
+    void getEventStreamMember(
+        @Mock ProtocolGenerator.GenerationContext context,
+        @Mock Model model,
+        @Mock StructureShape struct,
+        @Mock MemberShape eventStreamMember1,
+        @Mock ShapeId streamingMember1ShapeId,
+        @Mock UnionShape streamingTarget1
+    ) {
+        when(struct.members()).thenReturn(List.of(eventStreamMember1));
+        when(eventStreamMember1.getTarget()).thenReturn(streamingMember1ShapeId);
+        when(context.getModel()).thenReturn(model);
+        when(model.expectShape(streamingMember1ShapeId)).thenReturn(streamingTarget1);
+
+        when(streamingTarget1.hasTrait(StreamingTrait.class)).thenReturn(true);
+        when(streamingTarget1.isUnionShape()).thenReturn(true);
+        when(eventStreamMember1.hasTrait(StreamingTrait.class)).thenReturn(false);
+        when(eventStreamMember1.hasTrait(HttpPayloadTrait.class)).thenReturn(true);
+
+        MemberShape eventStreamMember = EventStreamGenerator.getEventStreamMember(
+            context,
+            struct
+        );
+
+        assertEquals(eventStreamMember1, eventStreamMember);
+    }
+
+    @Test
+    void getEventStreamMemberTooFew(
+        @Mock ProtocolGenerator.GenerationContext context,
+        @Mock StructureShape struct
+    ) {
+        when(struct.members()).thenReturn(List.of());
+        when(struct.getId()).thenReturn(ShapeId.from("namespace#Shape"));
+
+        try {
+            MemberShape eventStreamMember = EventStreamGenerator.getEventStreamMember(
+                context,
+                struct
+            );
+        } catch (CodegenException e) {
+            assertEquals(
+                "No event stream member found in " + struct.getId().toString(),
+                e.getMessage()
+            );
+        }
+    }
+
+    @Test
+    void getEventStreamMemberTooMany(
+        @Mock ProtocolGenerator.GenerationContext context,
+        @Mock Model model,
+        @Mock StructureShape struct,
+        @Mock MemberShape eventStreamMember1,
+        @Mock ShapeId streamingMember1ShapeId,
+        @Mock UnionShape streamingTarget1,
+        @Mock MemberShape eventStreamMember2,
+        @Mock ShapeId streamingMember2ShapeId,
+        @Mock UnionShape streamingTarget2
+    ) {
+        when(struct.members()).thenReturn(List.of(
+            eventStreamMember1,
+            eventStreamMember2
+        ));
+        when(context.getModel()).thenReturn(model);
+        when(struct.getId()).thenReturn(ShapeId.from("namespace#Shape"));
+
+        when(eventStreamMember1.getTarget()).thenReturn(streamingMember1ShapeId);
+        when(model.expectShape(streamingMember1ShapeId)).thenReturn(streamingTarget1);
+        when(streamingTarget1.hasTrait(StreamingTrait.class)).thenReturn(true);
+        when(streamingTarget1.isUnionShape()).thenReturn(true);
+        when(eventStreamMember1.hasTrait(StreamingTrait.class)).thenReturn(false);
+        when(eventStreamMember1.hasTrait(HttpPayloadTrait.class)).thenReturn(true);
+
+        when(eventStreamMember2.getTarget()).thenReturn(streamingMember2ShapeId);
+        when(model.expectShape(streamingMember2ShapeId)).thenReturn(streamingTarget2);
+        when(streamingTarget2.hasTrait(StreamingTrait.class)).thenReturn(true);
+        when(streamingTarget2.isUnionShape()).thenReturn(true);
+        when(eventStreamMember2.hasTrait(StreamingTrait.class)).thenReturn(false);
+        when(eventStreamMember2.hasTrait(HttpPayloadTrait.class)).thenReturn(true);
+
+        try {
+            MemberShape eventStreamMember = EventStreamGenerator.getEventStreamMember(
+                context,
+                struct
+            );
+        } catch (CodegenException e) {
+            assertEquals(
+                "More than one event stream member in " + struct.getId().toString(),
+                e.getMessage()
+            );
+        }
+    }
+}