fix: s3 custom treatment getbucketlocation response (#989)

Author: 0marperez

.changes/4197c56c-e17a-4fb3-a680-4ca384267d2b.json

@@ -0,0 +1,8 @@
+{
+    "id": "4197c56c-e17a-4fb3-a680-4ca384267d2b",
+    "type": "bugfix",
+    "description": "Correctly parse and handle `GetBucketLocation` responses",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#194"
+    ]
+}

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/s3/GetBucketLocationDeserializerIntegration.kt

@@ -0,0 +1,43 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization.s3
+
+import software.amazon.smithy.aws.traits.customizations.S3UnwrappedXmlOutputTrait
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.model.traits.UnwrappedXmlOutput
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.shapes.StructureShape
+import software.amazon.smithy.model.transform.ModelTransformer
+
+/**
+ * Applies the [UnwrappedXmlOutput] custom-made [annotation trait](https://smithy.io/2.0/spec/model.html?highlight=annotation#annotation-traits) to structures
+ * whose operation is annotated with `S3UnwrappedXmlOutput` trait to mark when special unwrapped xml output deserialization is required.
+ */
+class UnwrappedXmlOutputIntegration : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        model.expectShape<ServiceShape>(settings.service).isS3
+
+    override fun preprocessModel(model: Model, settings: KotlinSettings): Model {
+        val unwrappedStructures = model
+            .operationShapes
+            .filter { it.hasTrait<S3UnwrappedXmlOutputTrait>() }
+            .map { it.outputShape }
+            .toSet()
+
+        return ModelTransformer
+            .create()
+            .mapShapes(model) { shape ->
+                when {
+                    shape.id in unwrappedStructures ->
+                        (shape as StructureShape).toBuilder().addTrait(UnwrappedXmlOutput()).build()
+                    else -> shape
+                }
+            }
+    }
+}

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/s3/UnwrappedXmlOutputIntegration.kt

@@ -6,7 +6,6 @@ aws.sdk.kotlin.codegen.GradleGenerator
 aws.sdk.kotlin.codegen.AwsServiceConfigIntegration
 aws.sdk.kotlin.codegen.customization.s3.S3SigningConfig
 aws.sdk.kotlin.codegen.customization.s3.S3ErrorMetadataIntegration
-aws.sdk.kotlin.codegen.customization.s3.GetBucketLocationDeserializerIntegration
 aws.sdk.kotlin.codegen.customization.PresignableModelIntegration
 aws.sdk.kotlin.codegen.PresignerGenerator
 aws.sdk.kotlin.codegen.customization.apigateway.ApiGatewayAddAcceptHeader
@@ -31,3 +30,4 @@ aws.sdk.kotlin.codegen.customization.s3control.HostPrefixFilter
 aws.sdk.kotlin.codegen.customization.s3control.ClientConfigIntegration
 aws.sdk.kotlin.codegen.protocols.endpoints.BindAwsEndpointBuiltins
 aws.sdk.kotlin.codegen.customization.s3.HostPrefixRequestRouteFilter
+aws.sdk.kotlin.codegen.customization.s3.UnwrappedXmlOutputIntegration

codegen/smithy-aws-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration

@@ -0,0 +1,30 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization.s3
+
+import aws.sdk.kotlin.codegen.testutil.model
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.kotlin.codegen.test.defaultSettings
+import kotlin.test.assertTrue
+
+/**
+ * Verify [UnwrappedXmlOutputIntegration] is enabled for proper service (S3)
+ */
+class UnwrappedXmlOutputIntegrationTest {
+    @Test
+    fun testNonS3Model() {
+        val model = model("NotS3")
+        val actual = UnwrappedXmlOutputIntegration().enabledForService(model, model.defaultSettings())
+        assertFalse(actual)
+    }
+
+    @Test
+    fun testS3Model() {
+        val model = model("S3")
+        val actual = UnwrappedXmlOutputIntegration().enabledForService(model, model.defaultSettings())
+        assertTrue(actual)
+    }
+}

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/s3/UnwrappedXmlOutputIntegrationTest.kt

@@ -0,0 +1,66 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package aws.sdk.kotlin.services.s3.internal
+
+import aws.sdk.kotlin.services.s3.model.BucketLocationConstraint
+import aws.sdk.kotlin.services.s3.model.S3Exception
+import aws.sdk.kotlin.services.s3.transform.GetBucketLocationOperationDeserializer
+import aws.smithy.kotlin.runtime.http.Headers
+import aws.smithy.kotlin.runtime.http.HttpBody
+import aws.smithy.kotlin.runtime.http.HttpStatusCode
+import aws.smithy.kotlin.runtime.http.response.HttpResponse
+import aws.smithy.kotlin.runtime.operation.ExecutionContext
+import kotlinx.coroutines.runBlocking
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+
+class GetBucketLocationOperationDeserializerTest {
+    @Test
+    fun deserializeUnwrappedResponse() {
+        val responseXML = """
+            <?xml version="1.0" encoding="UTF-8"?>
+            <LocationConstraint xmlns="http://s3.amazonaws.com/doc/2006-03-01/">us-west-2</LocationConstraint>
+        """.trimIndent()
+
+        val response: HttpResponse = HttpResponse(
+            HttpStatusCode(200, "Success"),
+            Headers.invoke { },
+            HttpBody.fromBytes(responseXML.encodeToByteArray()),
+        )
+
+        val actual = runBlocking {
+            GetBucketLocationOperationDeserializer().deserialize(ExecutionContext(), response)
+        }
+
+        assertEquals(BucketLocationConstraint.UsWest2, actual.locationConstraint)
+    }
+
+    @Test
+    fun deserializeErrorMessage() {
+        val responseXML = """
+            <?xml version="1.0" encoding="UTF-8"?>
+            <Error>
+               <Message>Some message</Message>
+               <RequestId>Some request ID</RequestId>
+            </Error>
+        """.trimIndent()
+
+        val response: HttpResponse = HttpResponse(
+            HttpStatusCode(400, "Bad Request"),
+            Headers.invoke { },
+            HttpBody.fromBytes(responseXML.encodeToByteArray()),
+        )
+
+        val exception = assertThrows<S3Exception> {
+            runBlocking {
+                GetBucketLocationOperationDeserializer().deserialize(ExecutionContext(), response)
+            }
+        }
+
+        assertEquals("Some message", exception.message)
+    }
+}