Handle return type properly

tokio-macros/src/entry.rs

@@ -380,6 +380,23 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
         }
     };
 
+    let output_type = match &input.sig.output {
+        // For functions with no return value syn doesn't print anything,
+        // but that doesn't work as `Output` for our boxed `Future`, so
+        // default to `()` (the same type as the function output).
+        syn::ReturnType::Default => quote! { () },
+        syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
+    };
+
+    input.stmts.last_mut().map(|stmt| {
+        *stmt = quote! {
+            let _lst_stmt: #output_type = {
+                #stmt
+            };
+            _lst_stmt
+        };
+    });
+
     let body = input.body();
 
     // For test functions pin the body to the stack and use `Pin<&mut dyn
@@ -392,13 +409,6 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
     // We don't do this for the main function as it should only be used once so
     // there will be no benefit.
     let body = if is_test {
-        let output_type = match &input.sig.output {
-            // For functions with no return value syn doesn't print anything,
-            // but that doesn't work as `Output` for our boxed `Future`, so
-            // default to `()` (the same type as the function output).
-            syn::ReturnType::Default => quote! { () },
-            syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
-        };
         quote! {
             let body = async #body;
             #crate_path::pin!(body);