chore(codegen): read Endpoint Discovery Command from trait (#2524)

aws · Jun 29, 2021 · a8502a7 · a8502a7
1 parent 623f80f
commit a8502a7
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/...c/main/java/software/amazon/smithy/aws/typescript/codegen/AddEndpointDiscoveryPlugin.java b/...c/main/java/software/amazon/smithy/aws/typescript/codegen/AddEndpointDiscoveryPlugin.java
@@ -69,10 +69,9 @@ public List<RuntimeClientPlugin> getClientPlugins() {
                 RuntimeClientPlugin.builder()
                         .withConventions(AwsDependency.MIDDLEWARE_ENDPOINT_DISCOVERY.dependency,
                                 "EndpointDiscovery", RuntimeClientPlugin.Convention.HAS_CONFIG)
-                        // ToDo: The Endpoint Discovery Command Name needs to be read from ClientEndpointDiscoveryTrait.
                         .additionalResolveFunctionParamsSupplier((m, s, o) -> new HashMap<String, Object>() {{
                             put("endpointDiscoveryCommandCtor",
-                                    Symbol.builder().name("DescribeEndpointsCommand").build());
+                                    Symbol.builder().name(getClientDiscoveryCommand(s)).build());
                         }})
                         .servicePredicate((m, s) -> hasClientEndpointDiscovery(s))
                         .build(),
@@ -161,4 +160,13 @@ private static boolean isClientDiscoveredEndpointOptional(ServiceShape service,
         }
         return false;
     }
+
+    private static String getClientDiscoveryCommand(ServiceShape service) {
+        if (!hasClientEndpointDiscovery(service)) {
+            throw new CodegenException(
+                "EndpointDiscovery command discovery attempt for service without endpoint discovery"
+            );
+        }
+        return service.getTrait(ClientEndpointDiscoveryTrait.class).orElse(null).getOperation().getName() + "Command";
+    }
 }