Add Accept header tests

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

@@ -0,0 +1,103 @@
+/*
+ * 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")
+        println(actual)
+        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")
+        println(actual)
+        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")
+}