Add service builder test assertion

smithy-lang · Sep 10, 2022 · a19fdf6 · a19fdf6
1 parent ba8d772
commit a19fdf6
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/...azon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocolTestGenerator.kt b/...azon/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,