From e26d00ebaa5e158508944d69c40240284f93303b Mon Sep 17 00:00:00 2001
From: Martijn Gribnau <garm@ilumeo.com>
Date: Fri, 13 Oct 2023 00:09:38 +0200
Subject: [PATCH] feat: add option to use return type in tests

Especially useful for types which use the try operator, `?` i.e. Result or Option, such that you can use ? instead of unwrapping everything. Accepted by the #[test] macro
---
 parameterized-macro/src/generation.rs             |  3 ++-
 parameterized-macro/tests/cases.rs                |  3 +++
 parameterized-macro/tests/ok/21_result.rs         |  8 ++++++++
 parameterized-macro/tests/ok/22_try_operator.rs   |  9 +++++++++
 .../tests/ok/23_return_not_impl_try.rs            | 10 ++++++++++
 src/lib.rs                                        | 15 +++++++++++++++
 6 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 parameterized-macro/tests/ok/21_result.rs
 create mode 100644 parameterized-macro/tests/ok/22_try_operator.rs
 create mode 100644 parameterized-macro/tests/ok/23_return_not_impl_try.rs

diff --git a/parameterized-macro/src/generation.rs b/parameterized-macro/src/generation.rs
index 7a6b4f8..422a7a3 100644
--- a/parameterized-macro/src/generation.rs
+++ b/parameterized-macro/src/generation.rs
@@ -78,6 +78,7 @@ fn generate_test_case(
     let vis = &f.vis;
     let body_block = &f.block;
     let identifier = syn::Ident::new(&format!("case_{}", i), Span::call_site());
+    let return_type = &f.sig.output;
 
     // Construction let bindings for all parameters
     let bindings = parameters.iter().map(|(identifier, ty)| {
@@ -89,7 +90,7 @@ fn generate_test_case(
     quote::quote! {
         #[test]
         #(#attributes)*
-        #vis fn #identifier() {
+        #vis fn #identifier() #return_type {
             #(#bindings)*
 
             #body_block
diff --git a/parameterized-macro/tests/cases.rs b/parameterized-macro/tests/cases.rs
index 9ba59bf..bdc6c31 100644
--- a/parameterized-macro/tests/cases.rs
+++ b/parameterized-macro/tests/cases.rs
@@ -21,6 +21,9 @@ fn individual_cases() {
     t.pass("tests/ok/18_trailing_comma3.rs");
     t.pass("tests/ok/19_trailing_comma4.rs");
     t.pass("tests/ok/20_empty.rs");
+    t.pass("tests/ok/21_result.rs");
+    t.pass("tests/ok/22_try_operator.rs");
+    t.pass("tests/ok/23_return_not_impl_try.rs");
 
     t.compile_fail("tests/fail/id_already_defined.rs");
     t.compile_fail("tests/fail/inequal_amount_of_arg.rs");
diff --git a/parameterized-macro/tests/ok/21_result.rs b/parameterized-macro/tests/ok/21_result.rs
new file mode 100644
index 0000000..9b9768c
--- /dev/null
+++ b/parameterized-macro/tests/ok/21_result.rs
@@ -0,0 +1,8 @@
+use parameterized_macro::parameterized;
+
+#[parameterized(v = { 1, 2 })]
+fn my_test(v: i32) -> Result<(), ()> {
+    Ok(())
+}
+
+fn main() {}
diff --git a/parameterized-macro/tests/ok/22_try_operator.rs b/parameterized-macro/tests/ok/22_try_operator.rs
new file mode 100644
index 0000000..95c8516
--- /dev/null
+++ b/parameterized-macro/tests/ok/22_try_operator.rs
@@ -0,0 +1,9 @@
+use parameterized_macro::parameterized;
+
+#[parameterized(v = { 1, 2 })]
+fn my_test(v: i32) -> Result<(), ()> {
+    let unit = Err(())?;
+    Ok(unit)
+}
+
+fn main() {}
diff --git a/parameterized-macro/tests/ok/23_return_not_impl_try.rs b/parameterized-macro/tests/ok/23_return_not_impl_try.rs
new file mode 100644
index 0000000..f1c8b2e
--- /dev/null
+++ b/parameterized-macro/tests/ok/23_return_not_impl_try.rs
@@ -0,0 +1,10 @@
+use parameterized_macro::parameterized;
+
+struct NotTry;
+
+#[parameterized(_v = { 1, 2 })]
+fn my_test(_v: i32) -> NotTry {
+    NotTry
+}
+
+fn main() {}
diff --git a/src/lib.rs b/src/lib.rs
index 6f0109e..cfa6973 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -70,4 +70,19 @@ mod tests {
             input.unwrap()
         }
     }
+
+    mod result {
+        use super::*;
+
+        ide!();
+
+        #[pm(input = { 2, 3, 4 }, output = { 2, 3, 4 })]
+        fn ok(input: i32, output: i32) -> Result<(), ()> {
+            let result = Ok(input)?;
+
+            assert_eq!(result, output);
+
+            Ok(())
+        }
+    }
 }