From f1f9e8f80a4faffd5638b2885cb19772014e6a2e Mon Sep 17 00:00:00 2001
From: "Nate McMaster (AWS)" <88004173+mcmasn-amzn@users.noreply.github.com>
Date: Mon, 9 Dec 2024 13:54:30 -0800
Subject: [PATCH] fix: update test operation generator to find any operation in
 the service closure (#3922)

## Motivation and Context
When using operations nested under resources, the current search will
raise "NoSuchElementException". This fixes the search to include all
operations within the service closure

## Description
The error is visible in client generates like this:
```
Projection rust-client failed: java.util.NoSuchElementException: Collection contains no element matching the predicate.
java.util.NoSuchElementException: Collection contains no element matching the predicate.
        at software.amazon.smithy.rustsdk.endpoints.OperationInputTestGeneratorKt.operationId(OperationInputTestGenerator.kt:242)
        at software.amazon.smithy.rustsdk.endpoints.OperationInputTestGenerator$operationInvocation$1.invoke(OperationInputTestGenerator.kt:183)
        at software.amazon.smithy.rustsdk.endpoints.OperationInputTestGenerator$operationInvocation$1.invoke(OperationInputTestGenerator.kt:180)
```

## Testing
`gradle :aws:sdk-codegen:check`
## Checklist


- [ ] For changes to the smithy-rs codegen or runtime crates, I have
created a changelog entry Markdown file in the `.changelog` directory,
specifying "client," "server," or both in the `applies_to` key.
- [ ] For changes to the AWS SDK, generated SDK code, or SDK runtime
crates, I have created a changelog entry Markdown file in the
`.changelog` directory, specifying "aws-sdk-rust" in the `applies_to`
key.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
---
 .../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 =