petri/
test.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Infrastructure for defining tests.
5
6#[doc(hidden)]
7pub mod test_macro_support {
8    // UNSAFETY: Needed for linkme.
9    #![expect(unsafe_code)]
10
11    use super::TestCase;
12    pub use linkme;
13
14    #[linkme::distributed_slice]
15    pub static TESTS: [Option<fn() -> (&'static str, Vec<TestCase>)>];
16
17    // Always have at least one entry to work around linker bugs.
18    //
19    // See <https://github.com/llvm/llvm-project/issues/65855>.
20    #[linkme::distributed_slice(TESTS)]
21    static WORKAROUND: Option<fn() -> (&'static str, Vec<TestCase>)> = None;
22}
23
24use crate::PetriLogSource;
25use crate::TestArtifactRequirements;
26use crate::TestArtifacts;
27use crate::requirements::HostContext;
28use crate::requirements::TestCaseRequirements;
29use crate::requirements::can_run_test_with_context;
30use crate::tracing::try_init_tracing;
31use anyhow::Context as _;
32use petri_artifacts_core::ArtifactResolver;
33use std::panic::AssertUnwindSafe;
34use std::panic::catch_unwind;
35use test_macro_support::TESTS;
36
37/// Defines a single test from a value that implements [`RunTest`].
38#[macro_export]
39macro_rules! test {
40    ($f:ident, $req:expr) => {
41        $crate::multitest!(vec![
42            $crate::SimpleTest::new(stringify!($f), $req, $f, None,).into()
43        ]);
44    };
45}
46
47/// Defines a set of tests from a [`TestCase`].
48#[macro_export]
49macro_rules! multitest {
50    ($tests:expr) => {
51        const _: () = {
52            use $crate::test_macro_support::linkme;
53            #[linkme::distributed_slice($crate::test_macro_support::TESTS)]
54            #[linkme(crate = linkme)]
55            static TEST: Option<fn() -> (&'static str, Vec<$crate::TestCase>)> =
56                Some(|| (module_path!(), $tests));
57        };
58    };
59}
60
61/// A single test case.
62pub struct TestCase(Box<dyn DynRunTest>);
63
64impl TestCase {
65    /// Creates a new test case from a value that implements [`RunTest`].
66    pub fn new(test: impl 'static + RunTest) -> Self {
67        Self(Box::new(test))
68    }
69}
70
71impl<T: 'static + RunTest> From<T> for TestCase {
72    fn from(test: T) -> Self {
73        Self::new(test)
74    }
75}
76
77/// A single test, with module name.
78struct Test {
79    module: &'static str,
80    test: TestCase,
81    artifact_requirements: TestArtifactRequirements,
82}
83
84impl Test {
85    /// Returns all the tests defined in this crate.
86    fn all() -> impl Iterator<Item = Self> {
87        TESTS.iter().flatten().flat_map(|f| {
88            let (module, tests) = f();
89            tests.into_iter().filter_map(move |test| {
90                let mut artifact_requirements = test.0.artifact_requirements()?;
91                // All tests require the log directory.
92                artifact_requirements
93                    .require(petri_artifacts_common::artifacts::TEST_LOG_DIRECTORY);
94                Some(Self {
95                    module,
96                    artifact_requirements,
97                    test,
98                })
99            })
100        })
101    }
102
103    /// Returns the name of the test.
104    fn name(&self) -> String {
105        // Strip the crate name from the module path, for consistency with libtest.
106        match self.module.split_once("::") {
107            Some((_crate_name, rest)) => format!("{}::{}", rest, self.test.0.leaf_name()),
108            None => self.test.0.leaf_name().to_owned(),
109        }
110    }
111
112    fn run(
113        &self,
114        resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
115    ) -> anyhow::Result<()> {
116        let name = self.name();
117        let artifacts = resolve(&name, self.artifact_requirements.clone())
118            .context("failed to resolve artifacts")?;
119        let output_dir = artifacts.get(petri_artifacts_common::artifacts::TEST_LOG_DIRECTORY);
120        let logger = try_init_tracing(output_dir).context("failed to initialize tracing")?;
121
122        // Catch test panics in order to cleanly log the panic result. Without
123        // this, `libtest_mimic` will report the panic to stdout and fail the
124        // test, but the details won't end up in our per-test JSON log.
125        let r = catch_unwind(AssertUnwindSafe(|| {
126            self.test.0.run(
127                PetriTestParams {
128                    test_name: &name,
129                    logger: &logger,
130                },
131                &artifacts,
132            )
133        }));
134        let r = r.unwrap_or_else(|err| {
135            // The error from `catch_unwind` is almost always either a
136            // `&str` or a `String`, since that's what `panic!` produces.
137            let msg = err
138                .downcast_ref::<&str>()
139                .copied()
140                .or_else(|| err.downcast_ref::<String>().map(|x| x.as_str()));
141
142            let err = if let Some(msg) = msg {
143                anyhow::anyhow!("test panicked: {msg}")
144            } else {
145                anyhow::anyhow!("test panicked (unknown payload type)")
146            };
147            Err(err)
148        });
149        logger.log_test_result(&name, &r);
150        r
151    }
152
153    /// Returns a libtest-mimic trial to run the test.
154    fn trial(
155        self,
156        resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
157    ) -> libtest_mimic::Trial {
158        libtest_mimic::Trial::test(self.name(), move || {
159            self.run(resolve).map_err(|err| format!("{err:#}").into())
160        })
161    }
162}
163
164/// A test that can be run.
165///
166/// Register it to be run with [`test!`] or [`multitest!`].
167pub trait RunTest: Send {
168    /// The type of artifacts required by the test.
169    type Artifacts;
170
171    /// The leaf name of the test.
172    ///
173    /// To produce the full test name, this will be prefixed with the module
174    /// name where the test is defined.
175    fn leaf_name(&self) -> &str;
176    /// Returns the artifacts required by the test.
177    ///
178    /// Returns `None` if this test makes no sense for this host environment
179    /// (e.g., an x86_64 test on an aarch64 host) and should be left out of the
180    /// test list.
181    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts>;
182    /// Runs the test, which has been assigned `name`, with the given
183    /// `artifacts`.
184    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()>;
185    /// Returns the host requirements of the current test, if any.
186    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
187}
188
189trait DynRunTest: Send {
190    fn leaf_name(&self) -> &str;
191    fn artifact_requirements(&self) -> Option<TestArtifactRequirements>;
192    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()>;
193    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
194}
195
196impl<T: RunTest> DynRunTest for T {
197    fn leaf_name(&self) -> &str {
198        self.leaf_name()
199    }
200
201    fn artifact_requirements(&self) -> Option<TestArtifactRequirements> {
202        let mut requirements = TestArtifactRequirements::new();
203        self.resolve(&ArtifactResolver::collector(&mut requirements))?;
204        Some(requirements)
205    }
206
207    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()> {
208        let artifacts = self
209            .resolve(&ArtifactResolver::resolver(artifacts))
210            .context("test should have been skipped")?;
211        self.run(params, artifacts)
212    }
213
214    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
215        self.host_requirements()
216    }
217}
218
219/// Parameters passed to a [`RunTest`] when it is run.
220pub struct PetriTestParams<'a> {
221    /// The name of the running test.
222    pub test_name: &'a str,
223    /// The logger for the test.
224    pub logger: &'a PetriLogSource,
225}
226
227/// A test defined by an artifact resolver function and a run function.
228pub struct SimpleTest<A, F> {
229    leaf_name: &'static str,
230    resolve: A,
231    run: F,
232    /// Optional test requirements
233    pub host_requirements: Option<TestCaseRequirements>,
234}
235
236impl<A, AR, F, E> SimpleTest<A, F>
237where
238    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
239    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
240    E: Into<anyhow::Error>,
241{
242    /// Returns a new test with the given `leaf_name`, `resolve`, `run` functions,
243    /// and optional requirements.
244    pub fn new(
245        leaf_name: &'static str,
246        resolve: A,
247        run: F,
248        host_requirements: Option<TestCaseRequirements>,
249    ) -> Self {
250        SimpleTest {
251            leaf_name,
252            resolve,
253            run,
254            host_requirements,
255        }
256    }
257}
258
259impl<A, AR, F, E> RunTest for SimpleTest<A, F>
260where
261    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
262    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
263    E: Into<anyhow::Error>,
264{
265    type Artifacts = AR;
266
267    fn leaf_name(&self) -> &str {
268        self.leaf_name
269    }
270
271    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts> {
272        (self.resolve)(resolver)
273    }
274
275    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()> {
276        (self.run)(params, artifacts).map_err(Into::into)
277    }
278
279    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
280        self.host_requirements.as_ref()
281    }
282}
283
284#[derive(clap::Parser)]
285struct Options {
286    /// Lists the required artifacts for all tests.
287    #[clap(long)]
288    list_required_artifacts: bool,
289    #[clap(flatten)]
290    inner: libtest_mimic::Arguments,
291}
292
293/// Entry point for test binaries.
294pub fn test_main(
295    resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
296) -> ! {
297    let mut args = <Options as clap::Parser>::parse();
298    if args.list_required_artifacts {
299        // FUTURE: write this in a machine readable format.
300        for test in Test::all() {
301            println!("{}:", test.name());
302            for artifact in test.artifact_requirements.required_artifacts() {
303                println!("required: {artifact:?}");
304            }
305            for artifact in test.artifact_requirements.optional_artifacts() {
306                println!("optional: {artifact:?}");
307            }
308            println!();
309        }
310        std::process::exit(0);
311    }
312
313    // Always just use one thread to avoid interleaving logs and to avoid using
314    // too many resources. These tests are usually run under nextest, which will
315    // run them in parallel in separate processes with appropriate concurrency
316    // limits.
317    if !matches!(args.inner.test_threads, None | Some(1)) {
318        eprintln!("warning: ignoring value passed to --test-threads, using 1");
319    }
320    args.inner.test_threads = Some(1);
321
322    // Create the host context once to avoid repeated expensive queries
323    let host_context = futures::executor::block_on(HostContext::new());
324
325    let trials = Test::all()
326        .map(|test| {
327            let can_run = can_run_test_with_context(test.test.0.host_requirements(), &host_context);
328            test.trial(resolve).with_ignored_flag(!can_run)
329        })
330        .collect();
331
332    libtest_mimic::run(&args.inner, trials).exit();
333}