Update non-streaming error unmarshaller with the new mapping function (#6122)

* Update non-streaming error unmarshaller with the new mapping function

* Updating tests

Author: S-Saranya1

.changes/next-release/bugfix-AWSSDKforJavav2-b6c94e2.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Update non-streaming error unmarshalling to properly unmarshall exceptions to their expected types."
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java

@@ -39,6 +39,7 @@
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeType;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
 import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
 import software.amazon.awssdk.codegen.poet.client.traits.HttpChecksumRequiredTrait;
@@ -116,7 +117,6 @@ public MethodSpec initProtocolFactory(IntermediateModel model) {
             methodSpec.addCode("$L", hasAwsQueryCompatible());
         }
 
-        registerModeledExceptions(model, poetExtensions).forEach(methodSpec::addCode);
         methodSpec.addCode(";");
 
         return methodSpec.build();
@@ -170,11 +170,40 @@ public CodeBlock responseHandler(IntermediateModel model, OperationModel opModel
     public Optional<CodeBlock> errorResponseHandler(OperationModel opModel) {
         String protocolFactory = protocolFactoryLiteral(model, opModel);
 
-        return Optional.of(
-            CodeBlock.builder()
-                     .add("\n\n$T<$T> errorResponseHandler = createErrorResponseHandler($L, operationMetadata);",
-                          HttpResponseHandler.class, AwsServiceException.class, protocolFactory)
-                     .build());
+        CodeBlock.Builder builder = CodeBlock.builder();
+        ParameterizedTypeName metadataMapperType = ParameterizedTypeName.get(
+            ClassName.get(Function.class),
+            ClassName.get(String.class),
+            ParameterizedTypeName.get(Optional.class, ExceptionMetadata.class));
+
+        builder.add("\n$T exceptionMetadataMapper = errorCode -> {\n", metadataMapperType);
+        builder.add("if (errorCode == null) {\n");
+        builder.add("return $T.empty();\n", Optional.class);
+        builder.add("}\n");
+        builder.add("switch (errorCode) {\n");
+        model.getShapes().values().stream()
+             .filter(shape -> shape.getShapeType() == ShapeType.Exception)
+             .forEach(exceptionShape -> {
+                 String exceptionName = exceptionShape.getShapeName();
+                 String errorCode = exceptionShape.getErrorCode();
+
+                 builder.add("case $S:\n", errorCode);
+                 builder.add("return $T.of($T.builder()\n", Optional.class, ExceptionMetadata.class)
+                        .add(".errorCode($S)\n", errorCode);
+                 builder.add(populateHttpStatusCode(exceptionShape, model));
+                 builder.add(".exceptionBuilderSupplier($T::builder)\n",
+                             poetExtensions.getModelClassFromShape(exceptionShape))
+                        .add(".build());\n");
+             });
+
+        builder.add("default: return $T.empty();\n", Optional.class);
+        builder.add("}\n");
+        builder.add("};\n");
+
+        builder.add("$T<$T> errorResponseHandler = createErrorResponseHandler($L, operationMetadata, exceptionMetadataMapper);",
+                    HttpResponseHandler.class, AwsServiceException.class, protocolFactory);
+
+        return Optional.of(builder.build());
     }
 
     @Override
@@ -411,21 +440,6 @@ public Optional<MethodSpec> createErrorResponseHandler() {
         ClassName httpResponseHandler = ClassName.get(HttpResponseHandler.class);
         ClassName sdkBaseException = ClassName.get(AwsServiceException.class);
         TypeName responseHandlerOfException = ParameterizedTypeName.get(httpResponseHandler, sdkBaseException);
-
-        return Optional.of(MethodSpec.methodBuilder("createErrorResponseHandler")
-                                     .addParameter(BaseAwsJsonProtocolFactory.class, "protocolFactory")
-                                     .addParameter(JsonOperationMetadata.class, "operationMetadata")
-                                     .returns(responseHandlerOfException)
-                                     .addModifiers(Modifier.PRIVATE)
-                                     .addStatement("return protocolFactory.createErrorResponseHandler(operationMetadata)")
-                                     .build());
-    }
-
-    @Override
-    public Optional<MethodSpec> createEventstreamErrorResponseHandler() {
-        ClassName httpResponseHandler = ClassName.get(HttpResponseHandler.class);
-        ClassName sdkBaseException = ClassName.get(AwsServiceException.class);
-        TypeName responseHandlerOfException = ParameterizedTypeName.get(httpResponseHandler, sdkBaseException);
         ParameterizedTypeName mapperType = ParameterizedTypeName.get(ClassName.get(Function.class),
             ClassName.get(String.class), ParameterizedTypeName.get(Optional.class, ExceptionMetadata.class));