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, false).into()
43        ]);
44    };
45}
46
47/// Defines a single unstable test from a value that implements [`RunTest`].
48#[macro_export]
49macro_rules! unstable_test {
50    ($f:ident, $req:expr) => {
51        $crate::multitest!(vec![
52            $crate::SimpleTest::new(stringify!($f), $req, $f, None, true).into()
53        ]);
54    };
55}
56
57/// Defines a set of tests from a [`TestCase`].
58#[macro_export]
59macro_rules! multitest {
60    ($tests:expr) => {
61        const _: () = {
62            use $crate::test_macro_support::linkme;
63            #[linkme::distributed_slice($crate::test_macro_support::TESTS)]
64            #[linkme(crate = linkme)]
65            static TEST: Option<fn() -> (&'static str, Vec<$crate::TestCase>)> =
66                Some(|| (module_path!(), $tests));
67        };
68    };
69}
70
71/// A single test case.
72pub struct TestCase(Box<dyn DynRunTest>);
73
74impl TestCase {
75    /// Creates a new test case from a value that implements [`RunTest`].
76    pub fn new(test: impl 'static + RunTest) -> Self {
77        Self(Box::new(test))
78    }
79}
80
81impl<T: 'static + RunTest> From<T> for TestCase {
82    fn from(test: T) -> Self {
83        Self::new(test)
84    }
85}
86
87/// A single test, with module name.
88struct Test {
89    module: &'static str,
90    test: TestCase,
91    artifact_requirements: TestArtifactRequirements,
92}
93
94impl Test {
95    /// Returns all the tests defined in this crate.
96    fn all() -> impl Iterator<Item = Self> {
97        TESTS.iter().flatten().flat_map(|f| {
98            let (module, tests) = f();
99            tests.into_iter().filter_map(move |test| {
100                let mut artifact_requirements = test.0.artifact_requirements()?;
101                // All tests require the log directory.
102                artifact_requirements
103                    .require(petri_artifacts_common::artifacts::TEST_LOG_DIRECTORY);
104                Some(Self {
105                    module,
106                    artifact_requirements,
107                    test,
108                })
109            })
110        })
111    }
112
113    /// Returns the name of the test.
114    fn name(&self) -> String {
115        // Strip the crate name from the module path, for consistency with libtest.
116        match self.module.split_once("::") {
117            Some((_crate_name, rest)) => format!("{}::{}", rest, self.test.0.leaf_name()),
118            None => self.test.0.leaf_name().to_owned(),
119        }
120    }
121
122    fn run(
123        &self,
124        resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
125    ) -> anyhow::Result<()> {
126        let name = self.name();
127        let artifacts = resolve(&name, self.artifact_requirements.clone())
128            .context("failed to resolve artifacts")?;
129        let output_dir = artifacts.get(petri_artifacts_common::artifacts::TEST_LOG_DIRECTORY);
130        let logger = try_init_tracing(output_dir, tracing::level_filters::LevelFilter::DEBUG)
131            .context("failed to initialize tracing")?;
132        let mut post_test_hooks = Vec::new();
133
134        // Catch test panics in order to cleanly log the panic result. Without
135        // this, `libtest_mimic` will report the panic to stdout and fail the
136        // test, but the details won't end up in our per-test JSON log.
137        let r = catch_unwind(AssertUnwindSafe(|| {
138            self.test.0.run(
139                PetriTestParams {
140                    test_name: &name,
141                    logger: &logger,
142                    post_test_hooks: &mut post_test_hooks,
143                },
144                &artifacts,
145            )
146        }));
147        let r = r.unwrap_or_else(|err| {
148            // The error from `catch_unwind` is almost always either a
149            // `&str` or a `String`, since that's what `panic!` produces.
150            let msg = err
151                .downcast_ref::<&str>()
152                .copied()
153                .or_else(|| err.downcast_ref::<String>().map(|x| x.as_str()));
154
155            let err = if let Some(msg) = msg {
156                anyhow::anyhow!("test panicked: {msg}")
157            } else {
158                anyhow::anyhow!("test panicked (unknown payload type)")
159            };
160            Err(err)
161        });
162        logger.log_test_result(&name, &r, self.test.0.unstable());
163
164        for hook in post_test_hooks {
165            tracing::info!(name = hook.name(), "Running post-test hook");
166            if let Err(e) = hook.run(r.is_ok()) {
167                tracing::error!(
168                    error = e.as_ref() as &dyn std::error::Error,
169                    "Post-test hook failed"
170                );
171            } else {
172                tracing::info!("Post-test hook completed successfully");
173            }
174        }
175
176        r
177    }
178
179    /// Returns a libtest-mimic trial to run the test.
180    fn trial(
181        self,
182        resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
183    ) -> libtest_mimic::Trial {
184        libtest_mimic::Trial::test(self.name(), move || match self.run(resolve) {
185            Ok(()) => Ok(()),
186            Err(err)
187                if self.test.0.unstable()
188                    && std::env::var("PETRI_REPORT_UNSTABLE_FAIL")
189                        .ok()
190                        .is_none_or(|v| v.is_empty() || v == "0") =>
191            {
192                tracing::warn!("ignoring unstable test failure: {err:#}");
193                Ok(())
194            }
195            Err(err) => Err(format!("{err:#}").into()),
196        })
197    }
198}
199
200/// A test that can be run.
201///
202/// Register it to be run with [`test!`] or [`multitest!`].
203pub trait RunTest: Send {
204    /// The type of artifacts required by the test.
205    type Artifacts;
206
207    /// The leaf name of the test.
208    ///
209    /// To produce the full test name, this will be prefixed with the module
210    /// name where the test is defined.
211    fn leaf_name(&self) -> &str;
212    /// Returns the artifacts required by the test.
213    ///
214    /// Returns `None` if this test makes no sense for this host environment
215    /// (e.g., an x86_64 test on an aarch64 host) and should be left out of the
216    /// test list.
217    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts>;
218    /// Runs the test, which has been assigned `name`, with the given
219    /// `artifacts`.
220    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()>;
221    /// Returns the host requirements of the current test, if any.
222    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
223    /// Whether this test is unstable
224    fn unstable(&self) -> bool;
225}
226
227trait DynRunTest: Send {
228    fn leaf_name(&self) -> &str;
229    fn artifact_requirements(&self) -> Option<TestArtifactRequirements>;
230    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()>;
231    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
232    fn unstable(&self) -> bool;
233}
234
235impl<T: RunTest> DynRunTest for T {
236    fn leaf_name(&self) -> &str {
237        self.leaf_name()
238    }
239
240    fn artifact_requirements(&self) -> Option<TestArtifactRequirements> {
241        let mut requirements = TestArtifactRequirements::new();
242        self.resolve(&ArtifactResolver::collector(&mut requirements))?;
243        Some(requirements)
244    }
245
246    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()> {
247        let artifacts = self
248            .resolve(&ArtifactResolver::resolver(artifacts))
249            .context("test should have been skipped")?;
250        self.run(params, artifacts)
251    }
252
253    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
254        self.host_requirements()
255    }
256
257    fn unstable(&self) -> bool {
258        self.unstable()
259    }
260}
261
262/// Parameters passed to a [`RunTest`] when it is run.
263pub struct PetriTestParams<'a> {
264    /// The name of the running test.
265    pub test_name: &'a str,
266    /// The logger for the test.
267    pub logger: &'a PetriLogSource,
268    /// Any hooks that want to run after the test completes.
269    pub post_test_hooks: &'a mut Vec<PetriPostTestHook>,
270}
271
272/// A post-test hook to be run after the test completes, regardless of if it
273/// succeeds or fails.
274pub struct PetriPostTestHook {
275    /// The name of the hook.
276    name: String,
277    /// The hook function.
278    hook: Box<dyn FnOnce(bool) -> anyhow::Result<()>>,
279}
280
281impl PetriPostTestHook {
282    pub fn new(name: String, hook: impl FnOnce(bool) -> anyhow::Result<()> + 'static) -> Self {
283        Self {
284            name,
285            hook: Box::new(hook),
286        }
287    }
288
289    pub fn name(&self) -> &str {
290        &self.name
291    }
292
293    pub fn run(self, test_passed: bool) -> anyhow::Result<()> {
294        (self.hook)(test_passed)
295    }
296}
297
298/// A test defined by an artifact resolver function and a run function.
299pub struct SimpleTest<A, F> {
300    leaf_name: &'static str,
301    resolve: A,
302    run: F,
303    /// Optional test requirements
304    pub host_requirements: Option<TestCaseRequirements>,
305    unstable: bool,
306}
307
308impl<A, AR, F, E> SimpleTest<A, F>
309where
310    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
311    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
312    E: Into<anyhow::Error>,
313{
314    /// Returns a new test with the given `leaf_name`, `resolve`, `run` functions,
315    /// and optional requirements.
316    pub fn new(
317        leaf_name: &'static str,
318        resolve: A,
319        run: F,
320        host_requirements: Option<TestCaseRequirements>,
321        unstable: bool,
322    ) -> Self {
323        SimpleTest {
324            leaf_name,
325            resolve,
326            run,
327            host_requirements,
328            unstable,
329        }
330    }
331}
332
333impl<A, AR, F, E> RunTest for SimpleTest<A, F>
334where
335    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
336    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
337    E: Into<anyhow::Error>,
338{
339    type Artifacts = AR;
340
341    fn leaf_name(&self) -> &str {
342        self.leaf_name
343    }
344
345    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts> {
346        (self.resolve)(resolver)
347    }
348
349    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()> {
350        (self.run)(params, artifacts).map_err(Into::into)
351    }
352
353    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
354        self.host_requirements.as_ref()
355    }
356
357    fn unstable(&self) -> bool {
358        self.unstable
359    }
360}
361
362#[derive(clap::Parser)]
363struct Options {
364    /// Lists the required artifacts for all tests.
365    #[clap(long)]
366    list_required_artifacts: bool,
367    #[clap(flatten)]
368    inner: libtest_mimic::Arguments,
369}
370
371/// Entry point for test binaries.
372pub fn test_main(
373    resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
374) -> ! {
375    let mut args = <Options as clap::Parser>::parse();
376    if args.list_required_artifacts {
377        // FUTURE: write this in a machine readable format.
378        for test in Test::all() {
379            println!("{}:", test.name());
380            for artifact in test.artifact_requirements.required_artifacts() {
381                println!("required: {artifact:?}");
382            }
383            for artifact in test.artifact_requirements.optional_artifacts() {
384                println!("optional: {artifact:?}");
385            }
386            println!();
387        }
388        std::process::exit(0);
389    }
390
391    // Always just use one thread to avoid interleaving logs and to avoid using
392    // too many resources. These tests are usually run under nextest, which will
393    // run them in parallel in separate processes with appropriate concurrency
394    // limits.
395    if !matches!(args.inner.test_threads, None | Some(1)) {
396        eprintln!("warning: ignoring value passed to --test-threads, using 1");
397    }
398    args.inner.test_threads = Some(1);
399
400    // Create the host context once to avoid repeated expensive queries
401    let host_context = futures::executor::block_on(HostContext::new());
402
403    let trials = Test::all()
404        .map(|test| {
405            let can_run = can_run_test_with_context(test.test.0.host_requirements(), &host_context);
406            test.trial(resolve).with_ignored_flag(!can_run)
407        })
408        .collect();
409
410    libtest_mimic::run(&args.inner, trials).exit();
411}