Add nativelink_test macro for tests (#888)

This macro is to replace the tokio::test macro. It will allow us to
do some setup & cleanup on all tests without burdening the dev with
wrappers.

nativelink-macro/BUILD.bazel

@@ -0,0 +1,23 @@
+load(
+    "@rules_rust//rust:defs.bzl",
+    "rust_doc",
+    "rust_proc_macro",
+)
+
+rust_proc_macro(
+    name = "nativelink-macro",
+    srcs = [
+        "src/lib.rs",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@crates//:proc-macro2",
+        "@crates//:quote",
+        "@crates//:syn",
+    ],
+)
+
+rust_doc(
+    name = "docs",
+    crate = ":nativelink-macro",
+)

nativelink-macro/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "nativelink-macro"
+version = "0.3.0"
+edition = "2021"
+
+[lib]
+crate_type = ["proc-macro"]
+
+[dependencies]
+# TODO(allada) We currently need to pin these to specific version.
+# Some down-stream can't be upgraded just yet.
+proc-macro2 = "=1.0.78"
+quote = "=1.0.35"
+syn = "=2.0.52"

nativelink-macro/src/lib.rs

@@ -0,0 +1,41 @@
+// Copyright 2024 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+extern crate proc_macro;
+
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::{parse_macro_input, ItemFn};
+
+#[proc_macro_attribute]
+pub fn nativelink_test(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let attr = proc_macro2::TokenStream::from(attr);
+    let input_fn = parse_macro_input!(item as ItemFn);
+
+    let fn_name = &input_fn.sig.ident;
+    let fn_block = &input_fn.block;
+    let fn_inputs = &input_fn.sig.inputs;
+    let fn_output = &input_fn.sig.output;
+    let fn_attr = &input_fn.attrs;
+
+    let expanded = quote! {
+        #(#fn_attr)*
+        #[tokio::test(#attr)]
+        async fn #fn_name(#fn_inputs) #fn_output {
+            #fn_block
+        }
+    };
+
+    TokenStream::from(expanded)
+}

nativelink-scheduler/BUILD.bazel

@@ -61,6 +61,7 @@ rust_test_suite(
         "tests/utils/scheduler_utils.rs",
     ],
     proc_macro_deps = [
+        "//nativelink-macro",
         "@crates//:async-trait",
     ],
     deps = [

nativelink-scheduler/Cargo.toml

@@ -30,4 +30,6 @@ tonic = { version = "0.11.0", features = ["gzip", "tls"] }
 tracing = "0.1.40"
 
 [dev-dependencies]
+nativelink-macro = { path = "../nativelink-macro" }
+
 pretty_assertions = "1.4.0"

nativelink-scheduler/tests/action_messages_test.rs

@@ -17,6 +17,7 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
 use nativelink_error::Error;
+use nativelink_macro::nativelink_test;
 use nativelink_proto::build::bazel::remote::execution::v2::ExecuteResponse;
 use nativelink_proto::google::longrunning::{operation, Operation};
 use nativelink_proto::google::rpc::Status;
@@ -41,7 +42,7 @@ mod action_messages_tests {
 
     use super::*; // Must be declared in every module.
 
-    #[tokio::test]
+    #[nativelink_test]
     async fn action_state_any_url_test() -> Result<(), Error> {
         let action_state = ActionState {
             unique_qualifier: ActionInfoHashKey {
@@ -68,7 +69,7 @@ mod action_messages_tests {
         Ok(())
     }
 
-    #[tokio::test]
+    #[nativelink_test]
     async fn execute_response_status_message_is_some_on_success_test() -> Result<(), Error> {
         let execute_response: ExecuteResponse = ActionStage::Completed(ActionResult {
             output_files: vec![],
@@ -102,7 +103,7 @@ mod action_messages_tests {
         Ok(())
     }
 
-    #[tokio::test]
+    #[nativelink_test]
     async fn highest_priority_action_first() -> Result<(), Error> {
         const INSTANCE_NAME: &str = "foobar_instance_name";
 
@@ -158,7 +159,7 @@ mod action_messages_tests {
         Ok(())
     }
 
-    #[tokio::test]
+    #[nativelink_test]
     async fn equal_priority_earliest_first() -> Result<(), Error> {
         const INSTANCE_NAME: &str = "foobar_instance_name";