Make unitTest accept a closure using RustWriter (#1450)

* Make `unitTest` take `RustWriter.() -> Unit` * Convert `unitTests` which use dependencies to use `unitTest` with `RustWriter.() -> Unit`
smithy-lang · Jun 14, 2022 · 8911e86 · 8911e86
1 parent 66ff2bd
commit 8911e86
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 36 deletions.
diff --git a/...codegen/src/test/kotlin/software/amazon/smithy/rustsdk/EndpointConfigCustomizationTest.kt b/...codegen/src/test/kotlin/software/amazon/smithy/rustsdk/EndpointConfigCustomizationTest.kt
@@ -8,6 +8,8 @@ package software.amazon.smithy.rustsdk
 import org.junit.jupiter.api.Test
 import software.amazon.smithy.model.node.ObjectNode
 import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.rustlang.asType
+import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.smithy.CodegenVisitor
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
@@ -24,6 +26,10 @@ import software.amazon.smithy.rust.codegen.testutil.unitTest
 import software.amazon.smithy.rust.codegen.util.runCommand
 
 internal class EndpointConfigCustomizationTest {
+    private val codegenScope = arrayOf(
+        "http" to CargoDependency.Http.asType(),
+        "aws_types" to awsTypes(AwsTestRuntimeConfig).asType()
+    )
 
     private val model = """
         namespace test
@@ -163,21 +169,19 @@ internal class EndpointConfigCustomizationTest {
     fun `support region-specific endpoint overrides`() {
         validateEndpointCustomizationForService("test#TestService") { crate ->
             crate.lib {
-                it.addDependency(awsTypes(AwsTestRuntimeConfig))
-                it.addDependency(CargoDependency.Http)
-                it.unitTest(
-                    "region_override",
-                    """
-                    use aws_types::region::Region;
-                    use http::Uri;
-                    let conf = crate::config::Config::builder().build();
-                    let endpoint = conf.endpoint_resolver
-                        .resolve_endpoint(&Region::new("fips-ca-central-1")).expect("default resolver produces a valid endpoint");
-                    let mut uri = Uri::from_static("/?k=v");
-                    endpoint.set_endpoint(&mut uri, None);
-                    assert_eq!(uri, Uri::from_static("https://access-analyzer-fips.ca-central-1.amazonaws.com/?k=v"));
-                    """
-                )
+                it.unitTest("region_override") {
+                    rustTemplate(
+                        """
+                        let conf = crate::config::Config::builder().build();
+                        let endpoint = conf.endpoint_resolver
+                            .resolve_endpoint(&#{aws_types}::region::Region::new("fips-ca-central-1")).expect("default resolver produces a valid endpoint");
+                        let mut uri = #{http}::Uri::from_static("/?k=v");
+                        endpoint.set_endpoint(&mut uri, None);
+                        assert_eq!(uri, #{http}::Uri::from_static("https://access-analyzer-fips.ca-central-1.amazonaws.com/?k=v"));
+                        """,
+                        *codegenScope
+                    )
+                }
             }
         }
     }
@@ -186,27 +190,25 @@ internal class EndpointConfigCustomizationTest {
     fun `support region-agnostic services`() {
         validateEndpointCustomizationForService("test#NoRegions") { crate ->
             crate.lib {
-                it.addDependency(awsTypes(AwsTestRuntimeConfig))
-                it.addDependency(CargoDependency.Http)
-                it.unitTest(
-                    "global_services",
-                    """
-                    use aws_types::region::Region;
-                    use http::Uri;
-                    let conf = crate::config::Config::builder().build();
-                    let endpoint = conf.endpoint_resolver
-                        .resolve_endpoint(&Region::new("us-east-1")).expect("default resolver produces a valid endpoint");
-                    let mut uri = Uri::from_static("/?k=v");
-                    endpoint.set_endpoint(&mut uri, None);
-                    assert_eq!(uri, Uri::from_static("https://iam.amazonaws.com/?k=v"));
-
-                    let endpoint = conf.endpoint_resolver
-                        .resolve_endpoint(&Region::new("iam-fips")).expect("default resolver produces a valid endpoint");
-                    let mut uri = Uri::from_static("/?k=v");
-                    endpoint.set_endpoint(&mut uri, None);
-                    assert_eq!(uri, Uri::from_static("https://iam-fips.amazonaws.com/?k=v"));
-                    """
-                )
+                it.unitTest("global_services") {
+                    rustTemplate(
+                        """
+                        let conf = crate::config::Config::builder().build();
+                        let endpoint = conf.endpoint_resolver
+                            .resolve_endpoint(&#{aws_types}::region::Region::new("us-east-1")).expect("default resolver produces a valid endpoint");
+                        let mut uri = #{http}::Uri::from_static("/?k=v");
+                        endpoint.set_endpoint(&mut uri, None);
+                        assert_eq!(uri, #{http}::Uri::from_static("https://iam.amazonaws.com/?k=v"));
+
+                        let endpoint = conf.endpoint_resolver
+                            .resolve_endpoint(&#{aws_types}::region::Region::new("iam-fips")).expect("default resolver produces a valid endpoint");
+                        let mut uri = #{http}::Uri::from_static("/?k=v");
+                        endpoint.set_endpoint(&mut uri, None);
+                        assert_eq!(uri, #{http}::Uri::from_static("https://iam-fips.amazonaws.com/?k=v"));
+                        """,
+                        *codegenScope
+                    )
+                }
             }
         }
     }

diff --git a/codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/testutil/Rust.kt b/codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/testutil/Rust.kt
@@ -171,6 +171,18 @@ fun RustWriter.unitTest(
     }
 }
 
+/*
+ * Writes a Rust-style unit test
+ */
+fun RustWriter.unitTest(
+    name: String,
+    vararg args: Any,
+    block: RustWriter.() -> Unit
+): RustWriter {
+    raw("#[test]")
+    return rustBlock("fn $name()", *args, block = block)
+}
+
 /**
  * WriterDelegator used for test purposes
  *