From a19fdf6e9e2f236c543e269d6c82645bcd188147 Mon Sep 17 00:00:00 2001
From: Harry Barber <hlbarber@amazon.co.uk>
Date: Sat, 10 Sep 2022 08:01:07 +0000
Subject: [PATCH] Add service builder test assertion

---
 .../protocol/ServerProtocolTestGenerator.kt        | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocolTestGenerator.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocolTestGenerator.kt
index 221a4735c4..7bb0e3d89c 100644
--- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocolTestGenerator.kt
+++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocolTestGenerator.kt
@@ -106,6 +106,7 @@ class ServerProtocolTestGenerator(
         "SmithyHttp" to CargoDependency.SmithyHttp(coreCodegenContext.runtimeConfig).asType(),
         "Http" to CargoDependency.Http.asType(),
         "Hyper" to CargoDependency.Hyper.asType(),
+        "Tokio" to ServerCargoDependency.TokioDev.asType(),
         "Tower" to CargoDependency.Tower.asType(),
         "SmithyHttpServer" to ServerCargoDependency.SmithyHttpServer(coreCodegenContext.runtimeConfig).asType(),
         "AssertEq" to CargoDependency.PrettyAssertions.asType().member("assert_eq!"),
@@ -553,18 +554,25 @@ class ServerProtocolTestGenerator(
 
     /** Checks the request using the new service builder. */
     private fun checkRequest2(operationShape: OperationShape, operationSymbol: Symbol, httpRequestTestCase: HttpRequestTestCase, rustWriter: RustWriter) {
-        val (inputT, outputT) = operationInputOutputTypes[operationShape]!!
+        val (inputT, _) = operationInputOutputTypes[operationShape]!!
         val operationName = RustReservedWords.escapeIfNeeded(operationSymbol.name.toSnakeCase())
         rustWriter.rustTemplate(
             """
+            let (sender, mut receiver) = #{Tokio}::sync::mpsc::channel(1);
             let service = crate::service::$serviceName::unchecked_builder()
-                .$operationName(|input: $inputT| async move {
-                    #{Body:W}
+                .$operationName(move |input: $inputT| {
+                    let sender = sender.clone();
+                    async move {
+                        let result = { #{Body:W} };
+                        sender.send(()).await.expect("receiver dropped early");
+                        result
+                    }
                 })
                 .build::<#{Hyper}::body::Body>();
             let http_response = #{Tower}::ServiceExt::oneshot(service, http_request)
                 .await
                 .expect("unable to make an HTTP request");
+            assert!(receiver.recv().await.is_some())
             """,
             "Body" to checkRequestHandler(operationShape, httpRequestTestCase),
             *codegenScope,