Review comments

Author: dagnir

codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessor.java

@@ -22,9 +22,11 @@
 import java.util.Map;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.codegen.customization.CodegenCustomizationProcessor;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.model.service.ServiceModel;
 
 // TODO: Remove this when the long polling trait is formalized as a c2j trait.
@@ -48,6 +50,17 @@ public class LongPollingOperationProcessor implements CodegenCustomizationProces
         SERVICE_ID_TO_OPERATIONS_MAP = Collections.unmodifiableMap(serviceIdToOperationsMap);
     }
 
+    private final Map<String, List<String>> serviceIdToOperations;
+
+    public LongPollingOperationProcessor() {
+        this(SERVICE_ID_TO_OPERATIONS_MAP);
+    }
+
+    @SdkTestInternalApi
+    LongPollingOperationProcessor(Map<String, List<String>> serviceIdToOperations) {
+        this.serviceIdToOperations = serviceIdToOperations;
+    }
+
     @Override
     public void preprocess(ServiceModel serviceModel) {
         // no-op
@@ -56,15 +69,24 @@ public void preprocess(ServiceModel serviceModel) {
     @Override
     public void postprocess(IntermediateModel intermediateModel) {
         String serviceId = intermediateModel.getMetadata().getServiceId();
-        List<String> longPollingOperations = SERVICE_ID_TO_OPERATIONS_MAP.getOrDefault(serviceId, Collections.emptyList());
+
+        if (!serviceIdToOperations.containsKey(serviceId)) {
+            return;
+        }
+
+        if (intermediateModel.getMetadata().getProtocol() != Protocol.AWS_JSON) {
+            throw new IllegalArgumentException("Currently only AWS-JSON services can use the longPoll trait");
+        }
+
+        List<String> longPollingOperations = serviceIdToOperations.getOrDefault(serviceId, Collections.emptyList());
 
         for (String longPollingOperation : longPollingOperations) {
             OperationModel opModel = intermediateModel.getOperation(longPollingOperation);
             if (opModel != null) {
                 log.info("Setting the longPoll trait for {}#{}", serviceId, longPollingOperation);
                 opModel.setLongPolling(true);
             } else {
-                log.warn("Did not find operation {}#{}", serviceId, longPollingOperation);
+               throw new RuntimeException("Operation " + longPollingOperation + " not found for service " + serviceId);
             }
         }
     }

codegen/src/test/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessTest.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.codegen.customization.processors;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import software.amazon.awssdk.codegen.C2jModels;
+import software.amazon.awssdk.codegen.IntermediateModelBuilder;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
+import software.amazon.awssdk.codegen.model.service.ServiceMetadata;
+import software.amazon.awssdk.codegen.poet.ClientTestModels;
+
+public class LongPollingOperationProcessTest {
+
+    @ParameterizedTest
+    @MethodSource("nonJsonProtocols")
+    void postprocess_serviceInMap_serviceNotJson_throws(Protocol protocol) {
+        C2jModels c2jModels = ClientTestModels.awsJsonServiceC2jModels();
+
+        ServiceMetadata metadata = c2jModels.serviceModel().getMetadata();
+        metadata.setProtocols(Collections.singletonList(protocol.getValue()));
+
+        IntermediateModel intermediateModel = new IntermediateModelBuilder(c2jModels).build();
+
+        Map<String, List<String>> serviceToOperations = new HashMap<>();
+        serviceToOperations.put(metadata.getServiceId(), Collections.emptyList());
+        LongPollingOperationProcessor processor = new LongPollingOperationProcessor(serviceToOperations);
+
+        assertThatThrownBy(() -> processor.postprocess(intermediateModel))
+            .hasMessage("Currently only AWS-JSON services can use the longPoll trait");
+    }
+
+    @Test
+    void postprocess_operationNotFound_throws() {
+        IntermediateModel intermediateModel = ClientTestModels.awsJsonServiceModels();
+
+        Map<String, List<String>> serviceToOperations = new HashMap<>();
+        serviceToOperations.put(intermediateModel.getMetadata().getServiceId(), Collections.singletonList("SomeOperation"));
+        LongPollingOperationProcessor processor = new LongPollingOperationProcessor(serviceToOperations);
+
+        assertThatThrownBy(() -> processor.postprocess(intermediateModel))
+            .hasMessage("Operation SomeOperation not found for service Json Service");
+    }
+
+    private static Stream<Protocol> nonJsonProtocols() {
+        return Stream.of(Protocol.values()).filter(p -> p != Protocol.AWS_JSON);
+    }
+}