Fix eventstream legacy generation

codegen/src/main/java/software/amazon/awssdk/codegen/poet/eventstream/EventStreamVisitorBuilderImplSpec.java

@@ -46,7 +46,7 @@ public class EventStreamVisitorBuilderImplSpec extends EventStreamVisitorBuilder
     private final EventStreamSpecHelper eventStreamSpecHelper;
 
     public EventStreamVisitorBuilderImplSpec(GeneratorTaskParams params, OperationModel operationModel) {
-        super(params.getPoetExtensions(), operationModel);
+        super(params.getPoetExtensions(), params.getModel(), operationModel);
         this.intermediateModel = params.getModel();
         this.poetExt = params.getPoetExtensions();
         this.opModel = operationModel;
@@ -108,7 +108,7 @@ protected MethodSpec.Builder applyVisitSubTypeMethodSpecUpdates(TypeSpec.Builder
                                                                         MemberModel event) {
             ClassName eventSubType = poetExt.getModelClass(event.getShape().getShapeName());
             TypeName eventConsumerType = consumerType(eventSubType);
-            FieldSpec consumerField = FieldSpec.builder(eventConsumerType, eventConsumerName(event))
+            FieldSpec consumerField = FieldSpec.builder(eventConsumerType, eventStreamSpecHelper.eventConsumerName(event))
                                                .addModifiers(Modifier.PRIVATE, Modifier.FINAL)
                                                .build();
             typeBuilder.addField(consumerField);
@@ -138,7 +138,7 @@ protected MethodSpec.Builder applyOnSubTypeMethodSpecUpdates(TypeSpec.Builder ty
                                                                  MemberModel event) {
         ClassName eventSubType = poetExt.getModelClass(event.getShape().getShapeName());
         TypeName eventConsumerType = consumerType(eventSubType);
-        FieldSpec consumerField = FieldSpec.builder(eventConsumerType, eventConsumerName(event))
+        FieldSpec consumerField = FieldSpec.builder(eventConsumerType, eventStreamSpecHelper.eventConsumerName(event))
                                            .addModifiers(Modifier.PRIVATE)
                                            .build();
         typeBuilder.addField(consumerField);

codegen/src/main/java/software/amazon/awssdk/codegen/poet/eventstream/EventStreamVisitorBuilderInterfaceSpec.java

@@ -23,13 +23,14 @@
 import java.util.function.Consumer;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.codegen.docs.DocumentationBuilder;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
-import software.amazon.awssdk.utils.internal.CodegenNamingUtils;
+import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
 
 /**
  * Spec for builder interface for visitor.
@@ -39,12 +40,16 @@ public class EventStreamVisitorBuilderInterfaceSpec implements ClassSpec {
     private final OperationModel operationModel;
     private final ShapeModel eventStreamShape;
     private final ClassName visitorBuilderType;
+    private final EventStreamSpecHelper eventStreamSpecHelper;
 
-    EventStreamVisitorBuilderInterfaceSpec(PoetExtensions poetExt, OperationModel opModel) {
+    EventStreamVisitorBuilderInterfaceSpec(PoetExtensions poetExt,
+                                           IntermediateModel intermediateModel,
+                                           OperationModel opModel) {
         this.poetExt = poetExt;
         this.operationModel = opModel;
         this.eventStreamShape = EventStreamUtils.getEventStreamInResponse(operationModel.getOutputShape());
         this.visitorBuilderType = poetExt.eventStreamResponseHandlerVisitorBuilderType(opModel);
+        this.eventStreamSpecHelper = new EventStreamSpecHelper(eventStreamShape, intermediateModel);
     }
 
     @Override
@@ -121,13 +126,9 @@ protected MethodSpec.Builder applyOnSubTypeMethodSpecUpdates(TypeSpec.Builder ty
     private MethodSpec.Builder createOnSubTypeMethodSpec(MemberModel event) {
         ClassName eventSubType = poetExt.getModelClass(event.getShape().getShapeName());
         ParameterizedTypeName eventConsumerType = ParameterizedTypeName.get(ClassName.get(Consumer.class), eventSubType);
-        return MethodSpec.methodBuilder(eventConsumerName(event))
+        return MethodSpec.methodBuilder(eventStreamSpecHelper.eventConsumerName(event))
                          .addModifiers(Modifier.PUBLIC)
                          .addParameter(ParameterSpec.builder(eventConsumerType, "c").build())
                          .returns(visitorBuilderType);
     }
-
-    protected final String eventConsumerName(MemberModel event) {
-        return "on" + CodegenNamingUtils.pascalCase(event.getName());
-    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/eventstream/EventStreamVisitorInterfaceSpec.java

@@ -35,13 +35,15 @@
  */
 class EventStreamVisitorInterfaceSpec implements ClassSpec {
     private final PoetExtensions poetExt;
+    private final IntermediateModel intermediateModel;
     private final OperationModel operationModel;
     private final ShapeModel eventStreamShape;
     private final ClassName eventStreamBaseClass;
     private final EventStreamSpecHelper eventStreamSpecHelper;
 
     EventStreamVisitorInterfaceSpec(IntermediateModel intermediateModel, PoetExtensions poetExt, OperationModel operationModel) {
         this.poetExt = poetExt;
+        this.intermediateModel = intermediateModel;
         this.operationModel = operationModel;
         this.eventStreamShape = EventStreamUtils.getEventStreamInResponse(operationModel.getOutputShape());
         this.eventStreamBaseClass = poetExt.getModelClassFromShape(eventStreamShape);
@@ -72,7 +74,8 @@ protected TypeSpec.Builder createTypeSpec() {
                         .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                         .addJavadoc("Visitor for subtypes of {@link $T}.", eventStreamBaseClass)
                         .addMethod(createBuilderMethodSpec())
-                        .addType(new EventStreamVisitorBuilderInterfaceSpec(poetExt, operationModel).poetSpec());
+                        .addType(new EventStreamVisitorBuilderInterfaceSpec(poetExt, intermediateModel, operationModel)
+                                .poetSpec());
     }
 
     protected TypeSpec.Builder finalizeTypeSpec(TypeSpec.Builder builder) {

codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/EventStreamSpecHelper.java

@@ -96,4 +96,11 @@ public String eventTypeEnumValue(MemberModel eventModel) {
     public String eventBuilderMethodName(MemberModel eventModel) {
         return String.format("%sBuilder", StringUtils.uncapitalize(eventModel.getName()));
     }
+
+    public String eventConsumerName(MemberModel eventModel) {
+        if (useLegacyGenerationScheme(eventModel)) {
+            return "on" + eventModel.getShape().getShapeName();
+        }
+        return "on" + eventModel.getName();
+    }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/json/customization.config

@@ -15,6 +15,6 @@
     "createMethodParams": ["param1", "param2", "param3"]
   },
   "useLegacyEventGenerationScheme": {
-    "EventStream": ["EventOne", "event-two"]
+    "EventStream": ["EventOne", "event-two", "eventThree"]
   }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/json/service-2.json

@@ -417,6 +417,10 @@
         },
         "secondEventOne": {
           "shape": "EventOne"
+        },
+        // Legacy generation uses shape name == event name
+        "eventThree": {
+          "shape": "LegacyEventThree"
         }
       },
       "eventstream": true
@@ -439,6 +443,15 @@
       },
       "event": true
     },
+    "LegacyEventThree": {
+      "type": "structure",
+      "members": {
+        "Baz": {
+          "shape": "String"
+        }
+      },
+      "event": true
+    },
     "EventStreamOperationWithOnlyInputRequest": {
       "type": "structure",
       "required": [

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-async-client-class.java

@@ -305,6 +305,7 @@ public CompletableFuture<Void> eventStreamOperation(EventStreamOperationRequest
                     EventStreamTaggedUnionPojoSupplier.builder().putSdkPojoSupplier("EventOne", EventStream::eventOneBuilder)
                             .putSdkPojoSupplier("EventTheSecond", EventStream::eventTheSecondBuilder)
                             .putSdkPojoSupplier("secondEventOne", EventStream::secondEventOneBuilder)
+                            .putSdkPojoSupplier("eventThree", EventStream::eventThreeBuilder)
                             .defaultSdkPojoSupplier(() -> new SdkPojoBuilder(EventStream.UNKNOWN)).build());
 
             HttpResponseHandler<AwsServiceException> errorResponseHandler = createErrorResponseHandler(protocolFactory,
@@ -468,6 +469,7 @@ public CompletableFuture<Void> eventStreamOperationWithOnlyOutput(
                     EventStreamTaggedUnionPojoSupplier.builder().putSdkPojoSupplier("EventOne", EventStream::eventOneBuilder)
                             .putSdkPojoSupplier("EventTheSecond", EventStream::eventTheSecondBuilder)
                             .putSdkPojoSupplier("secondEventOne", EventStream::secondEventOneBuilder)
+                            .putSdkPojoSupplier("eventThree", EventStream::eventThreeBuilder)
                             .defaultSdkPojoSupplier(() -> new SdkPojoBuilder(EventStream.UNKNOWN)).build());
 
             HttpResponseHandler<AwsServiceException> errorResponseHandler = createErrorResponseHandler(protocolFactory,

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/eventstream/test-response-handler.java

@@ -107,6 +107,17 @@ default void visitSecondEventOne(EventOne event) {
             visitDefault(event);
         }
 
+        /**
+         * Invoked when a {@link LegacyEventThree} is encountered. If this is not overridden, the event will be given to
+         * {@link #visitDefault(EventStream)}.
+         *
+         * @param event
+         *        Event being visited
+         */
+        default void visit(LegacyEventThree event) {
+            visitDefault(event);
+        }
+
         /**
          * Builder for {@link Visitor}. The {@link Visitor} class may also be extended for a more traditional style but
          * this builder allows for a more functional way of creating a visitor will callback methods.
@@ -153,6 +164,15 @@ interface Builder {
              * @return This builder for method chaining.
              */
             Builder onSecondEventOne(Consumer<EventOne> c);
+
+            /**
+             * Callback to invoke when a {@link LegacyEventThree} is visited.
+             *
+             * @param c
+             *        Callback to process the event.
+             * @return This builder for method chaining.
+             */
+            Builder onLegacyEventThree(Consumer<LegacyEventThree> c);
         }
     }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/eventstream/test-visitor-builder.java

@@ -15,6 +15,8 @@ final class DefaultEventStreamOperationVisitorBuilder implements EventStreamOper
 
     private Consumer<EventOne> onSecondEventOne;
 
+    private Consumer<LegacyEventThree> onLegacyEventThree;
+
     @Override
     public EventStreamOperationResponseHandler.Visitor.Builder onDefault(Consumer<EventStream> c) {
         this.onDefault = c;
@@ -44,6 +46,12 @@ public EventStreamOperationResponseHandler.Visitor.Builder onSecondEventOne(Cons
         return this;
     }
 
+    @Override
+    public EventStreamOperationResponseHandler.Visitor.Builder onLegacyEventThree(Consumer<LegacyEventThree> c) {
+        this.onLegacyEventThree = c;
+        return this;
+    }
+
     @Generated("software.amazon.awssdk:codegen")
     static class VisitorFromBuilder implements EventStreamOperationResponseHandler.Visitor {
         private final Consumer<EventStream> onDefault;
@@ -54,6 +62,8 @@ static class VisitorFromBuilder implements EventStreamOperationResponseHandler.V
 
         private final Consumer<EventOne> onSecondEventOne;
 
+        private final Consumer<LegacyEventThree> onLegacyEventThree;
+
         VisitorFromBuilder(DefaultEventStreamOperationVisitorBuilder builder) {
             this.onDefault = builder.onDefault != null ? builder.onDefault
                     : EventStreamOperationResponseHandler.Visitor.super::visitDefault;
@@ -63,6 +73,8 @@ static class VisitorFromBuilder implements EventStreamOperationResponseHandler.V
                     : EventStreamOperationResponseHandler.Visitor.super::visitEventTheSecond;
             this.onSecondEventOne = builder.onSecondEventOne != null ? builder.onSecondEventOne
                     : EventStreamOperationResponseHandler.Visitor.super::visitSecondEventOne;
+            this.onLegacyEventThree = builder.onLegacyEventThree != null ? builder.onLegacyEventThree
+                    : EventStreamOperationResponseHandler.Visitor.super::visit;
         }
 
         @Override
@@ -84,6 +96,11 @@ public void visitEventTheSecond(EventTwo event) {
         public void visitSecondEventOne(EventOne event) {
             onSecondEventOne.accept(event);
         }
+
+        @Override
+        public void visit(LegacyEventThree event) {
+            onLegacyEventThree.accept(event);
+        }
     }
 }
 

test/codegen-generated-classes-test/src/main/resources/codegen-resources/eventstreams/service-2.json

@@ -75,7 +75,7 @@
     "EventStream": {
       "type": "structure",
       "members": {
-        "EventOne": {
+        "TheEventOne": {
           "shape": "EventOne"
         },
         "LegacyGeneratedEvent": {

test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/eventstreams/EventDispatchTest.java

@@ -47,7 +47,7 @@ public class EventDispatchTest {
     private Consumer<EventStream> onDefaultConsumer;
 
     @Mock
-    private Consumer<EventOne> eventOneConsumer;
+    private Consumer<EventOne> theEventOneConsumer;
 
     @Mock
     private Consumer<EventOne> legacyGeneratedEventConsumer;
@@ -62,28 +62,28 @@ public class EventDispatchTest {
     public ExpectedException expected = ExpectedException.none();
 
     @Test
-    public void test_acceptEventOne_correctVisitorMethodCalled() {
-        EventStream eventStream = EventStream.eventOneBuilder().build();
+    public void test_acceptTheEventOne_correctVisitorMethodCalled() {
+        EventStream eventStream = EventStream.theEventOneBuilder().build();
         eventStream.accept(visitor);
 
-        verify(visitor).visitEventOne(Mockito.eq((EventOne) eventStream));
+        verify(visitor).visitTheEventOne(Mockito.eq((EventOne) eventStream));
         verifyNoMoreInteractions(visitor);
     }
 
     @Test
-    public void test_acceptEventOne_visitorBuiltWithBuilder_correctVisitorMethodCalled() {
+    public void test_acceptTheEventOne_visitorBuiltWithBuilder_correctVisitorMethodCalled() {
         EventStreamOperationResponseHandler.Visitor visitor = visitorBuiltWithBuilder();
-        EventStream eventStream = EventStream.eventOneBuilder().build();
+        EventStream eventStream = EventStream.theEventOneBuilder().build();
 
         eventStream.accept(visitor);
 
-        verify(eventOneConsumer).accept(eq((EventOne) eventStream));
+        verify(theEventOneConsumer).accept(eq((EventOne) eventStream));
         verifyNoMoreConsumerInteractions();
     }
 
     @Test
     public void test_acceptLegacyGeneratedEvent_correctVisitorMethodCalled() {
-        EventStream eventStream = EventStream.legacyGeneratedEventBuilder().build();
+        EventStream eventStream = EventOne.builder().build();
         eventStream.accept(visitor);
 
         // Note: notice the visit() method rather than visitLegacyGeneratedEvent()
@@ -94,7 +94,7 @@ public void test_acceptLegacyGeneratedEvent_correctVisitorMethodCalled() {
     @Test
     public void test_acceptLegacyGeneratedEvent_visitorBuiltWithBuilder_correctVisitorMethodCalled() {
         EventStreamOperationResponseHandler.Visitor visitor = visitorBuiltWithBuilder();
-        EventStream eventStream = EventStream.legacyGeneratedEventBuilder().build();
+        EventStream eventStream = EventOne.builder().build();
 
         eventStream.accept(visitor);
 
@@ -155,15 +155,15 @@ public void test_acceptOnBaseClass_UnCustomizedEvent_throwsException() {
     private EventStreamOperationResponseHandler.Visitor visitorBuiltWithBuilder() {
         return EventStreamOperationResponseHandler.Visitor.builder()
                 .onDefault(onDefaultConsumer)
-                .onEventOne(eventOneConsumer)
+                .onTheEventOne(theEventOneConsumer)
                 .onEventTwo(eventTwoConsumer)
-                .onLegacyGeneratedEvent(legacyGeneratedEventConsumer)
+                .onEventOne(legacyGeneratedEventConsumer)
                 .onSecondEventTwo(secondEventTwoConsumer)
                 .build();
     }
 
     private void verifyNoMoreConsumerInteractions() {
-        verifyNoMoreInteractions(onDefaultConsumer, eventOneConsumer, eventTwoConsumer, legacyGeneratedEventConsumer,
+        verifyNoMoreInteractions(onDefaultConsumer, theEventOneConsumer, eventTwoConsumer, legacyGeneratedEventConsumer,
                 secondEventTwoConsumer);
     }
 }