fix: account ID based endpoints

Author: 0marperez

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/KotlinJmespathExpressionVisitor.kt

@@ -19,13 +19,11 @@ import software.amazon.smithy.kotlin.codegen.model.isEnum
 import software.amazon.smithy.kotlin.codegen.model.targetOrSelf
 import software.amazon.smithy.kotlin.codegen.model.traits.OperationInput
 import software.amazon.smithy.kotlin.codegen.model.traits.OperationOutput
+import software.amazon.smithy.kotlin.codegen.utils.doubleQuote
 import software.amazon.smithy.kotlin.codegen.utils.dq
 import software.amazon.smithy.kotlin.codegen.utils.getOrNull
 import software.amazon.smithy.kotlin.codegen.utils.toCamelCase
-import software.amazon.smithy.model.shapes.ListShape
-import software.amazon.smithy.model.shapes.MapShape
-import software.amazon.smithy.model.shapes.MemberShape
-import software.amazon.smithy.model.shapes.Shape
+import software.amazon.smithy.model.shapes.*
 
 private val suffixSequence = sequenceOf("") + generateSequence(2) { it + 1 }.map(Int::toString) // "", "2", "3", etc.
 
@@ -526,11 +524,20 @@ class KotlinJmespathExpressionVisitor(
             ".$expr"
         }
 
-    private fun VisitedExpression.getKeys(): String {
-        val keys = this.shape?.targetOrSelf(ctx.model)?.allMembers
-            ?.keys?.joinToString(", ", "listOf(", ")") { "\"$it\"" }
-        return keys ?: "listOf<String>()"
-    }
+    // TODO: Support maps as objects throughout the visitor
+    // This visitor assumes JMESPath 'objects' will be classes. DDB assumes JMESPath 'objects' will be maps
+    private fun VisitedExpression.getKeys(): String =
+        if (this.shape?.targetOrSelf(ctx.model)?.type == ShapeType.MAP) {
+            "${this.identifier}?.keys?.toList()"
+        } else {
+            this
+                .shape
+                ?.targetOrSelf(ctx.model)
+                ?.allMembers
+                ?.keys
+                ?.joinToString(", ", "listOf(", ")") { it.doubleQuote() }
+                ?: "listOf<String>()"
+        }
 
     private fun VisitedExpression.getValues(): String {
         val values = this.shape?.targetOrSelf(ctx.model)?.allMembers?.keys

codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/endpoints/OperationContextParamsTest.kt

@@ -88,7 +88,7 @@ class OperationContextParamsTest {
     }
 
     @Test
-    fun testKeysFunctionPath() {
+    fun testKeysFunctionPathWithStructure() {
         val input = """
             structure TestOperationRequest {
                 Object: Object
@@ -114,4 +114,31 @@ class OperationContextParamsTest {
 
         codegen(pathResultType, path, input).shouldContainOnlyOnceWithDiff(expected)
     }
+
+    @Test
+    fun testKeysFunctionPathWithMap() {
+        val input = """
+            structure TestOperationRequest {
+                Object: StringMap
+            }
+            
+            map StringMap {
+                key: String
+                value: String
+            }
+        """.trimIndent()
+
+        val path = "keys(Object)"
+        val pathResultType = "stringArray"
+
+        val expected = """
+            @Suppress("UNCHECKED_CAST")
+            val input = request.context[HttpOperationContext.OperationInput] as TestOperationRequest
+            val object = input.object
+            val keys = object?.keys?.toList()
+            builder.foo = keys
+        """.formatForTest("    ")
+
+        codegen(pathResultType, path, input).shouldContainOnlyOnceWithDiff(expected)
+    }
 }

tests/codegen/waiter-tests/model/function-keys.smithy

@@ -18,14 +18,14 @@ use smithy.waiters#waitable
             }
         ]
     },
-    KeysFunctionPrimitivesIntegerEquals: {
+    KeysFunctionMapStringEquals: {
         acceptors: [
             {
                 state: "success",
                 matcher: {
                     output: {
-                        path: "keys(primitives)",
-                        expected: "integer",
+                        path: "keys(maps.strings)",
+                        expected: "key",
                         comparator: "anyStringEquals"
                     }
                 }

tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionKeysTest.kt

@@ -5,11 +5,12 @@
 
 package com.test
 
+import com.test.model.EntityMaps
 import com.test.model.EntityPrimitives
 import com.test.model.GetFunctionKeysEqualsRequest
 import com.test.model.GetFunctionKeysEqualsResponse
 import com.test.utils.successTest
-import com.test.waiters.waitUntilKeysFunctionPrimitivesIntegerEquals
+import com.test.waiters.waitUntilKeysFunctionMapStringEquals
 import com.test.waiters.waitUntilKeysFunctionPrimitivesStringEquals
 import kotlin.test.Test
 
@@ -22,9 +23,9 @@ class FunctionKeysTest {
     )
 
     @Test
-    fun testKeysFunctionPrimitivesIntegerEquals() = successTest(
+    fun testKeysFunctionMapStringEquals() = successTest(
         GetFunctionKeysEqualsRequest { name = "test" },
-        WaitersTestClient::waitUntilKeysFunctionPrimitivesIntegerEquals,
-        GetFunctionKeysEqualsResponse { primitives = EntityPrimitives { } },
+        WaitersTestClient::waitUntilKeysFunctionMapStringEquals,
+        GetFunctionKeysEqualsResponse { maps = EntityMaps { strings = mapOf("key" to "value") } },
     )
 }