misc: add Accept header to all RpcV2Cbor requests (#1170)

Author: lauzadis

.changes/2daf7e93-a9e4-4c6b-9c0c-b083a1864c09.json

@@ -0,0 +1,5 @@
+{
+    "id": "2daf7e93-a9e4-4c6b-9c0c-b083a1864c09",
+    "type": "misc",
+    "description": "Add `Accept` header to all RpcV2Cbor requests"
+}

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/aws/protocols/RpcV2Cbor.kt

@@ -26,6 +26,9 @@ import software.amazon.smithy.model.traits.TimestampFormatTrait
 import software.amazon.smithy.model.traits.UnitTypeTrait
 import software.amazon.smithy.protocol.traits.Rpcv2CborTrait
 
+private const val ACCEPT_HEADER = "application/cbor"
+private const val ACCEPT_HEADER_EVENT_STREAM = "application/vnd.amazon.eventstream"
+
 class RpcV2Cbor : AwsHttpBindingProtocolGenerator() {
     override val protocol: ShapeId = Rpcv2CborTrait.ID
 
@@ -59,13 +62,16 @@ class RpcV2Cbor : AwsHttpBindingProtocolGenerator() {
             }
         }
 
-        // Requests with event stream responses MUST include an `Accept` header set to the value `application/vnd.amazon.eventstream`
-        val eventStreamsAcceptHeaderMiddleware = object : ProtocolMiddleware {
-            private val mutateHeadersMiddleware = MutateHeadersMiddleware(extraHeaders = mapOf("Accept" to "application/vnd.amazon.eventstream"))
-
-            override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean = op.isOutputEventStream(ctx.model)
-            override val name: String = "RpcV2CborEventStreamsAcceptHeaderMiddleware"
-            override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) = mutateHeadersMiddleware.render(ctx, op, writer)
+        // Add `Accept` header with value `application/cbor` for standard responses
+        // and `application/vnd.amazon.eventstream` for event stream responses
+        val acceptHeaderMiddleware = object : ProtocolMiddleware {
+            override val name: String = "RpcV2CborAcceptHeaderMiddleware"
+            override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+                val acceptHeaderValue = if (op.isOutputEventStream(ctx.model)) ACCEPT_HEADER_EVENT_STREAM else ACCEPT_HEADER
+                MutateHeadersMiddleware(
+                    extraHeaders = mapOf("Accept" to acceptHeaderValue),
+                ).render(ctx, op, writer)
+            }
         }
 
         // Emit a metric to track usage of RpcV2Cbor
@@ -79,7 +85,7 @@ class RpcV2Cbor : AwsHttpBindingProtocolGenerator() {
         return super.getDefaultHttpMiddleware(ctx) + listOf(
             smithyProtocolHeaderMiddleware,
             validateSmithyProtocolHeaderMiddleware,
-            eventStreamsAcceptHeaderMiddleware,
+            acceptHeaderMiddleware,
             businessMetricsMiddleware,
         )
     }

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/aws/protocols/RpcV2CborTest.kt

@@ -0,0 +1,101 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.kotlin.codegen.aws.protocols
+
+import software.amazon.smithy.kotlin.codegen.test.*
+import kotlin.test.Test
+
+class RpcV2CborTest {
+    val model = """
+        ${"$"}version: "2"
+
+        namespace com.test
+        
+        use smithy.protocols#rpcv2Cbor
+        use aws.api#service
+
+        @rpcv2Cbor
+        @service(sdkId: "CborExample")
+        service CborExample {
+            version: "1.0.0",
+            operations: [GetFoo, GetFooStreaming]
+        }
+
+        @http(method: "POST", uri: "/foo")
+        operation GetFoo {}
+        
+        @http(method: "POST", uri: "/foo-streaming")
+        operation GetFooStreaming {
+            input := {}
+            output := {
+                events: FooEvents
+            }
+        }
+        
+        // Model taken from https://smithy.io/2.0/spec/streaming.html#event-streams
+        @streaming
+        union FooEvents {
+            up: Movement
+            down: Movement
+            left: Movement
+            right: Movement
+            throttlingError: ThrottlingError
+        }
+        
+        structure Movement {
+           velocity: Float
+        }
+        
+        @error("client")
+        @retryable(throttling: true)
+        structure ThrottlingError {}
+        """.toSmithyModel()
+
+    @Test
+    fun testStandardAcceptHeader() {
+        val ctx = model.newTestContext("CborExample")
+
+        val generator = RpcV2Cbor()
+        generator.generateProtocolClient(ctx.generationCtx)
+
+        ctx.generationCtx.delegator.finalize()
+        ctx.generationCtx.delegator.flushWriters()
+
+        val actual = ctx.manifest.expectFileString("/src/main/kotlin/com/test/DefaultTestClient.kt")
+        val getFooMethod = actual.lines("    override suspend fun getFoo(input: GetFooRequest): GetFooResponse {", "    }")
+
+        val expectedHeaderMutation = """
+        op.install(
+            MutateHeaders().apply {
+                append("Accept", "application/cbor")
+            }
+        )
+        """.replaceIndent("        ")
+        getFooMethod.shouldContainOnlyOnceWithDiff(expectedHeaderMutation)
+    }
+
+    @Test
+    fun testEventStreamAcceptHeader() {
+        val ctx = model.newTestContext("CborExample")
+
+        val generator = RpcV2Cbor()
+        generator.generateProtocolClient(ctx.generationCtx)
+
+        ctx.generationCtx.delegator.finalize()
+        ctx.generationCtx.delegator.flushWriters()
+
+        val actual = ctx.manifest.expectFileString("/src/main/kotlin/com/test/DefaultTestClient.kt")
+        val getFooMethod = actual.lines("    override suspend fun <T> getFooStreaming(input: GetFooStreamingRequest, block: suspend (GetFooStreamingResponse) -> T): T {", "    }")
+
+        val expectedHeaderMutation = """
+        op.install(
+            MutateHeaders().apply {
+                append("Accept", "application/vnd.amazon.eventstream")
+            }
+        )
+        """.replaceIndent("        ")
+        getFooMethod.shouldContainOnlyOnceWithDiff(expectedHeaderMutation)
+    }
+}

codegen/smithy-kotlin-codegen-testutils/src/main/kotlin/software/amazon/smithy/kotlin/codegen/test/CodegenTestUtils.kt

@@ -28,6 +28,7 @@ import software.amazon.smithy.model.traits.TimestampFormatTrait
 import software.amazon.smithy.model.traits.Trait
 import software.amazon.smithy.protocol.traits.Rpcv2CborTrait
 import software.amazon.smithy.utils.StringUtils
+import kotlin.test.assertNotEquals
 
 // This file houses test classes and functions relating to the code generator (protocols, serializers, etc)
 // Items contained here should be relatively high-level, utilizing all members of codegen classes, Smithy, and
@@ -274,3 +275,19 @@ fun KotlinCodegenPlugin.Companion.createSymbolProvider(
  * create a new [KotlinWriter] using the test context package name
  */
 fun TestContext.newWriter(): KotlinWriter = KotlinWriter(generationCtx.settings.pkg.name)
+
+/**
+ * Get all the lines between [fromLine] and [toLine]
+ */
+fun String.lines(fromLine: String, toLine: String): String {
+    val allLines = lines()
+
+    val fromIdx = allLines.indexOf(fromLine)
+    assertNotEquals(-1, fromIdx, """Could not find from line "$fromLine" in all lines""")
+
+    val toIdxOffset = allLines.drop(fromIdx + 1).indexOf(toLine)
+    assertNotEquals(-1, toIdxOffset, """Could not find to line "$toLine" in all lines""")
+
+    val toIdx = toIdxOffset + fromIdx + 1
+    return allLines.subList(fromIdx, toIdx + 1).joinToString("\n")
+}

runtime/serde/serde-xml/common/test/aws/smithy/kotlin/runtime/serde/xml/XmlStreamReaderTest.kt

@@ -108,7 +108,6 @@ class XmlStreamReaderTest {
         """.trimIndent().encodeToByteArray()
 
         val actual = xmlStreamReader(payload).allTokens()
-        println(actual)
 
         assertEquals(6, actual.size)
         assertIs<XmlToken.BeginElement>(actual.first())