apollo_proc_macros: timed test macros

Author: victorkstarkware

Cargo.lock

@@ -544,3 +544,280 @@ pub fn log_every_n_sec(input: TokenStream) -> TokenStream {
 
     TokenStream::from(expanded)
 }
+
+// Constant to control the default time limit
+const DEFAULT_TIME_LIMIT_MS: u64 = 200;
+
+/// A timed version of the `#[test]` attribute that fails if the test exceeds a time limit.
+///
+/// # Usage
+///
+/// ```rust,ignore
+/// #[timed_test]
+/// fn my_test() {
+///     // test code
+/// }
+///
+/// // With custom time limit (in milliseconds)
+/// #[timed_test(500)]
+/// fn my_slower_test() {
+///     // test code
+/// }
+/// ```
+///
+/// The default time limit is 200ms. If the test exceeds this limit, it will panic with a message
+/// indicating how long it took.
+#[proc_macro_attribute]
+pub fn timed_test(attr: TokenStream, input: TokenStream) -> TokenStream {
+    let time_limit_ms = parse_time_limit(attr, DEFAULT_TIME_LIMIT_MS);
+    let input_fn = parse_macro_input!(input as ItemFn);
+
+    let fn_name = &input_fn.sig.ident;
+    let fn_attrs = &input_fn.attrs;
+    let fn_vis = &input_fn.vis;
+    let fn_sig = &input_fn.sig;
+    let fn_block = &input_fn.block;
+
+    let expanded = quote! {
+        #[test]
+        #(#fn_attrs)*
+        #fn_vis #fn_sig {
+            let start = std::time::Instant::now();
+            #fn_block
+            let elapsed = start.elapsed();
+            let elapsed_ms = elapsed.as_millis() as u64;
+            if elapsed_ms > #time_limit_ms {
+                panic!(
+                    "Test `{}` exceeded time limit: took {}ms, limit is {}ms",
+                    stringify!(#fn_name),
+                    elapsed_ms,
+                    #time_limit_ms
+                );
+            }
+        }
+    };
+
+    TokenStream::from(expanded)
+}
+
+/// A timed version of the `#[tokio::test]` attribute that fails if the test exceeds a time limit.
+///
+/// # Usage
+///
+/// ```rust,ignore
+/// #[timed_tokio_test]
+/// async fn my_test() {
+///     // test code
+/// }
+///
+/// // With custom time limit (in milliseconds)
+/// #[timed_tokio_test(1000)]
+/// async fn my_slower_test() {
+///     // test code
+/// }
+/// ```
+///
+/// The default time limit is 200ms. If the test exceeds this limit, it will panic with a message
+/// indicating how long it took.
+#[proc_macro_attribute]
+pub fn timed_tokio_test(attr: TokenStream, input: TokenStream) -> TokenStream {
+    let time_limit_ms = parse_time_limit(attr, DEFAULT_TIME_LIMIT_MS);
+    let input_fn = parse_macro_input!(input as ItemFn);
+
+    let fn_name = &input_fn.sig.ident;
+    let fn_attrs = &input_fn.attrs;
+    let fn_vis = &input_fn.vis;
+    let fn_sig = &input_fn.sig;
+    let fn_block = &input_fn.block;
+
+    let time_limit_duration = quote! {
+        std::time::Duration::from_millis(#time_limit_ms)
+    };
+
+    let expanded = quote! {
+        #[tokio::test]
+        #(#fn_attrs)*
+        #fn_vis #fn_sig {
+            let start = std::time::Instant::now();
+            let test_future = async move {
+                #fn_block
+            };
+            match tokio::time::timeout(#time_limit_duration, test_future).await {
+                Ok(_) => {
+                    let elapsed = start.elapsed();
+                    let elapsed_ms = elapsed.as_millis() as u64;
+                    // Even if we completed within timeout, check if we exceeded the limit
+                    // (timeout might have been slightly longer due to scheduling)
+                    if elapsed_ms > #time_limit_ms {
+                        panic!(
+                            "Test `{}` exceeded time limit: took {}ms, limit is {}ms",
+                            stringify!(#fn_name),
+                            elapsed_ms,
+                            #time_limit_ms
+                        );
+                    }
+                }
+                Err(_) => {
+                    let elapsed = start.elapsed();
+                    let elapsed_ms = elapsed.as_millis() as u64;
+                    panic!(
+                        "Test `{}` exceeded time limit: took {}ms, limit is {}ms",
+                        stringify!(#fn_name),
+                        elapsed_ms,
+                        #time_limit_ms
+                    );
+                }
+            }
+        }
+    };
+
+    TokenStream::from(expanded)
+}
+
+/// A timed version of the `#[rstest]` attribute that fails if the test exceeds a time limit.
+///
+/// # Usage
+///
+/// ```rust,ignore
+/// #[timed_rstest]
+/// fn my_test(param: u32) {
+///     // test code
+/// }
+///
+/// // With custom time limit (in milliseconds)
+/// #[timed_rstest(500)]
+/// fn my_slower_test(param: u32) {
+///     // test code
+/// }
+///
+/// // Can be combined with tokio::test for async tests
+/// #[timed_rstest]
+/// #[tokio::test]
+/// async fn my_async_test(param: u32) {
+///     // test code
+/// }
+/// ```
+///
+/// The default time limit is 200ms. If the test exceeds this limit, it will panic with a message
+/// indicating how long it took.
+///
+/// Note: When combined with `#[tokio::test]`, the timing will work correctly for async tests.
+#[proc_macro_attribute]
+pub fn timed_rstest(attr: TokenStream, input: TokenStream) -> TokenStream {
+    let time_limit_ms = parse_time_limit(attr, DEFAULT_TIME_LIMIT_MS);
+    let input_fn = parse_macro_input!(input as ItemFn);
+
+    let fn_name = &input_fn.sig.ident;
+    let fn_attrs = &input_fn.attrs;
+    let fn_vis = &input_fn.vis;
+    let fn_sig = &input_fn.sig;
+    let fn_block = &input_fn.block;
+
+    // Check if this is an async function by looking at the signature
+    let is_async = fn_sig.asyncness.is_some();
+
+    // Check if tokio::test is in the attributes
+    // #[tokio::test] has path segments: ["tokio", "test"]
+    let has_tokio_test = fn_attrs.iter().any(|attr| {
+        let path = attr.path();
+        if path.segments.len() == 2 {
+            path.segments[0].ident == "tokio" && path.segments[1].ident == "test"
+        } else {
+            false
+        }
+    });
+
+    let is_async_test = is_async || has_tokio_test;
+
+    if is_async_test {
+        let time_limit_duration = quote! {
+            std::time::Duration::from_millis(#time_limit_ms)
+        };
+
+        // Separate tokio::test from other attributes
+        let (tokio_test_attrs, other_attrs): (Vec<_>, Vec<_>) = fn_attrs.iter().partition(|attr| {
+            let path = attr.path();
+            path.segments.len() == 2
+                && path.segments[0].ident == "tokio"
+                && path.segments[1].ident == "test"
+        });
+
+        let expanded = quote! {
+            #(#tokio_test_attrs)*
+            #[::rstest::rstest]
+            #(#other_attrs)*
+            #fn_vis #fn_sig {
+                let start = std::time::Instant::now();
+                let result = tokio::time::timeout(#time_limit_duration, async {
+                    #fn_block
+                }).await;
+                match result {
+                    Ok(_) => {
+                        let elapsed = start.elapsed();
+                        let elapsed_ms = elapsed.as_millis() as u64;
+                        if elapsed_ms > #time_limit_ms {
+                            panic!(
+                                "Test `{}` exceeded time limit: took {}ms, limit is {}ms",
+                                stringify!(#fn_name),
+                                elapsed_ms,
+                                #time_limit_ms
+                            );
+                        }
+                    }
+                    Err(_) => {
+                        let elapsed = start.elapsed();
+                        let elapsed_ms = elapsed.as_millis() as u64;
+                        panic!(
+                            "Test `{}` exceeded time limit: took {}ms, limit is {}ms",
+                            stringify!(#fn_name),
+                            elapsed_ms,
+                            #time_limit_ms
+                        );
+                    }
+                }
+            }
+        };
+
+        TokenStream::from(expanded)
+    } else {
+        let expanded = quote! {
+            #[::rstest::rstest]
+            #(#fn_attrs)*
+            #fn_vis #fn_sig {
+                let start = std::time::Instant::now();
+                #fn_block
+                let elapsed = start.elapsed();
+                let elapsed_ms = elapsed.as_millis() as u64;
+                if elapsed_ms > #time_limit_ms {
+                    panic!(
+                        "Test `{}` exceeded time limit: took {}ms, limit is {}ms",
+                        stringify!(#fn_name),
+                        elapsed_ms,
+                        #time_limit_ms
+                    );
+                }
+            }
+        };
+
+        TokenStream::from(expanded)
+    }
+}
+
+/// Helper function to parse the time limit from the attribute.
+/// Returns the default if the attribute is empty.
+fn parse_time_limit(attr: TokenStream, default: u64) -> u64 {
+    if attr.is_empty() {
+        return default;
+    }
+
+    let attr_str = attr.to_string();
+    let trimmed = attr_str.trim();
+
+    if trimmed.is_empty() {
+        return default;
+    }
+
+    trimmed.parse::<u64>().unwrap_or_else(|_| {
+        panic!("Expected a positive integer for time limit in milliseconds, got: {}", trimmed)
+    })
+}

crates/apollo_proc_macros/src/lib.rs

@@ -16,6 +16,7 @@ metrics-exporter-prometheus.workspace = true
 papyrus_common.workspace = true
 prometheus-parse.workspace = true
 rstest.workspace = true
+tokio.workspace = true
 
 [lints]
 workspace = true

crates/apollo_proc_macros_tests/Cargo.toml

@@ -0,0 +1,70 @@
+use std::thread;
+use std::time::Duration;
+
+use apollo_proc_macros::{timed_rstest, timed_test, timed_tokio_test};
+
+// Test basic timed_test macro
+#[timed_test]
+fn test_fast_sync_test() {
+    // This should complete well within 200ms
+    assert_eq!(1 + 1, 2);
+}
+
+#[timed_test(50)]
+fn test_custom_time_limit() {
+    // This should complete well within 50ms
+    thread::sleep(Duration::from_millis(10));
+}
+
+// Test timed_tokio_test macro
+#[timed_tokio_test]
+async fn test_fast_async_test() {
+    // This should complete well within 200ms
+    tokio::time::sleep(Duration::from_millis(10)).await;
+    assert_eq!(1 + 1, 2);
+}
+
+#[timed_tokio_test(100)]
+async fn test_custom_async_time_limit() {
+    // This should complete well within 100ms
+    tokio::time::sleep(Duration::from_millis(20)).await;
+}
+
+// Test timed_rstest macro
+#[timed_rstest]
+#[case(1)]
+#[case(2)]
+fn test_rstest_sync(#[case] value: u32) {
+    // This should complete well within 200ms
+    assert!(value > 0);
+}
+
+#[timed_rstest(100)]
+#[case(1)]
+#[case(2)]
+fn test_rstest_sync_custom_limit(#[case] value: u32) {
+    // This should complete well within 100ms
+    thread::sleep(Duration::from_millis(10));
+    assert!(value > 0);
+}
+
+// Test timed_rstest with tokio::test
+#[timed_rstest]
+#[tokio::test]
+#[case(1)]
+#[case(2)]
+async fn test_rstest_async(#[case] value: u32) {
+    // This should complete well within 200ms
+    tokio::time::sleep(Duration::from_millis(10)).await;
+    assert!(value > 0);
+}
+
+#[timed_rstest(150)]
+#[tokio::test]
+#[case(1)]
+#[case(2)]
+async fn test_rstest_async_custom_limit(#[case] value: u32) {
+    // This should complete well within 150ms
+    tokio::time::sleep(Duration::from_millis(20)).await;
+    assert!(value > 0);
+}