From d31a1805daa1e44b9f91fcbf9230a56feadcf11b Mon Sep 17 00:00:00 2001
From: Michal Rostecki <vadorovsky@protonmail.com>
Date: Mon, 20 Mar 2023 12:41:54 +0100
Subject: [PATCH] integration-test: Add `tokio_integration_test` macro

This new macro runs a test in a Tokio runtime and it can be used for
asynchronous tests (defined as `async fn`).
---
 test/integration-test-macros/Cargo.toml |  1 +
 test/integration-test-macros/src/lib.rs | 29 ++++++++++++++++++++++++-
 test/integration-test/Cargo.toml        |  2 ++
 test/integration-test/src/tests/mod.rs  |  3 ++-
 4 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/test/integration-test-macros/Cargo.toml b/test/integration-test-macros/Cargo.toml
index e28cdc5c..29788ccb 100644
--- a/test/integration-test-macros/Cargo.toml
+++ b/test/integration-test-macros/Cargo.toml
@@ -6,6 +6,7 @@ publish = false
 
 [dependencies]
 quote = "1"
+proc-macro2 = "1.0"
 syn = {version = "1.0", features = ["full"]}
 
 [lib]
diff --git a/test/integration-test-macros/src/lib.rs b/test/integration-test-macros/src/lib.rs
index 297159ed..9818b7f5 100644
--- a/test/integration-test-macros/src/lib.rs
+++ b/test/integration-test-macros/src/lib.rs
@@ -1,6 +1,7 @@
 use proc_macro::TokenStream;
+use proc_macro2::Span;
 use quote::quote;
-use syn::{parse_macro_input, ItemFn};
+use syn::{parse_macro_input, Ident, ItemFn};
 
 #[proc_macro_attribute]
 pub fn integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
@@ -17,3 +18,29 @@ pub fn integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
     };
     TokenStream::from(expanded)
 }
+
+#[proc_macro_attribute]
+pub fn tokio_integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let item = parse_macro_input!(item as ItemFn);
+    let name = &item.sig.ident;
+    let name_str = &item.sig.ident.to_string();
+    let sync_name_str = format!("sync_{name_str}");
+    let sync_name = Ident::new(&sync_name_str, Span::call_site());
+    let expanded = quote! {
+        #item
+
+        fn #sync_name() {
+            let rt = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+            rt.block_on(#name());
+        }
+
+        inventory::submit!(crate::IntegrationTest {
+            name: concat!(module_path!(), "::", #sync_name_str),
+            test_fn: #sync_name,
+        });
+    };
+    TokenStream::from(expanded)
+}
diff --git a/test/integration-test/Cargo.toml b/test/integration-test/Cargo.toml
index 6b02b4a8..c8f7ca02 100644
--- a/test/integration-test/Cargo.toml
+++ b/test/integration-test/Cargo.toml
@@ -10,6 +10,7 @@ aya = { path = "../../aya" }
 aya-obj = { path = "../../aya-obj" }
 clap = { version = "4", features = ["derive"] }
 env_logger = "0.10"
+futures-core = "0.3"
 inventory = "0.3"
 integration-test-macros = { path = "../integration-test-macros" }
 lazy_static = "1"
@@ -20,3 +21,4 @@ rbpf = "0.1.0"
 regex = "1"
 tempfile = "3.3.0"
 libtest-mimic = "0.6.0"
+tokio = { version = "1.24", features = ["rt", "rt-multi-thread"] }
diff --git a/test/integration-test/src/tests/mod.rs b/test/integration-test/src/tests/mod.rs
index 127b037d..a3a21c46 100644
--- a/test/integration-test/src/tests/mod.rs
+++ b/test/integration-test/src/tests/mod.rs
@@ -11,7 +11,8 @@ pub mod rbpf;
 pub mod relocations;
 pub mod smoke;
 
-pub use integration_test_macros::integration_test;
+pub use integration_test_macros::{integration_test, tokio_integration_test};
+
 #[derive(Debug)]
 pub struct IntegrationTest {
     pub name: &'static str,