From a17b5c7465ee36df686fb38f741efca7bb5879ef Mon Sep 17 00:00:00 2001
From: Nate McMaster <mcmasn@amazon.com>
Date: Tue, 26 Nov 2024 20:18:44 +0000
Subject: [PATCH] fix: update test operation generator to find any operation in
 the service closure

When using operations nested under resources, the current search will raise "NoSuchElementException".
This fixes the search to include all operations within the service closure
---
 .../endpoints/OperationInputTestGenerator.kt  |  6 ++-
 .../OperationInputTestGeneratorTests.kt       | 44 +++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/endpoints/OperationInputTestGenerator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/endpoints/OperationInputTestGenerator.kt
index 489aef1d4e..4fc455c061 100644
--- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/endpoints/OperationInputTestGenerator.kt
+++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/endpoints/OperationInputTestGenerator.kt
@@ -5,6 +5,7 @@
 
 package software.amazon.smithy.rustsdk.endpoints
 
+import software.amazon.smithy.model.knowledge.TopDownIndex
 import software.amazon.smithy.model.node.Node
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.ShapeId
@@ -222,4 +223,7 @@ class OperationInputTestGenerator(_ctx: ClientCodegenContext, private val test:
 }
 
 fun ClientCodegenContext.operationId(testOperationInput: EndpointTestOperationInput): ShapeId =
-    this.serviceShape.allOperations.first { it.name == testOperationInput.operationName }
+    TopDownIndex.of(this.model)
+        .getContainedOperations(this.serviceShape)
+        .map { it.toShapeId() }
+        .first { it.name == testOperationInput.operationName }
diff --git a/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/OperationInputTestGeneratorTests.kt b/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/OperationInputTestGeneratorTests.kt
index 472d815995..f38addd230 100644
--- a/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/OperationInputTestGeneratorTests.kt
+++ b/aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/OperationInputTestGeneratorTests.kt
@@ -55,6 +55,50 @@ class OperationInputTestGeneratorTests {
         assertEquals("operations#Ping", operationId.toString())
     }
 
+    @Test
+    fun `finds operation shape by name from nested operations`() {
+        val prefix = "\$version: \"2\""
+        val operationModel =
+            """
+            $prefix
+            namespace operations.bells
+
+            resource Bell {
+                operations: [Ding]
+            }
+
+            operation Ding {}
+            """.trimIndent()
+        val serviceModel =
+            """
+            $prefix
+            namespace service
+
+            use operations.bells#Bell
+
+            service MyService {
+                resources: [Bell]
+            }
+            """.trimIndent()
+
+        val model =
+            Model.assembler()
+                .discoverModels()
+                .addUnparsedModel("operation.smithy", operationModel)
+                .addUnparsedModel("main.smithy", serviceModel)
+                .assemble()
+                .unwrap()
+
+        val context = testClientCodegenContext(model)
+        val testOperationInput =
+            EndpointTestOperationInput.builder()
+                .operationName("Ding")
+                .build()
+
+        val operationId = context.operationId(testOperationInput)
+        assertEquals("operations.bells#Ding", operationId.toString())
+    }
+
     @Test
     fun `fails for operation name not found`() {
         val model =