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        let mut post_test_hooks = Vec::new();
122
123        // Catch test panics in order to cleanly log the panic result. Without
124        // this, `libtest_mimic` will report the panic to stdout and fail the
125        // test, but the details won't end up in our per-test JSON log.
126        let r = catch_unwind(AssertUnwindSafe(|| {
127            self.test.0.run(
128                PetriTestParams {
129                    test_name: &name,
130                    logger: &logger,
131                    post_test_hooks: &mut post_test_hooks,
132                },
133                &artifacts,
134            )
135        }));
136        let r = r.unwrap_or_else(|err| {
137            // The error from `catch_unwind` is almost always either a
138            // `&str` or a `String`, since that's what `panic!` produces.
139            let msg = err
140                .downcast_ref::<&str>()
141                .copied()
142                .or_else(|| err.downcast_ref::<String>().map(|x| x.as_str()));
143
144            let err = if let Some(msg) = msg {
145                anyhow::anyhow!("test panicked: {msg}")
146            } else {
147                anyhow::anyhow!("test panicked (unknown payload type)")
148            };
149            Err(err)
150        });
151        logger.log_test_result(&name, &r);
152
153        for hook in post_test_hooks {
154            tracing::info!(name = hook.name(), "Running post-test hook");
155            if let Err(e) = hook.run(r.is_ok()) {
156                tracing::error!(
157                    error = e.as_ref() as &dyn std::error::Error,
158                    "Post-test hook failed"
159                );
160            } else {
161                tracing::info!("Post-test hook completed successfully");
162            }
163        }
164
165        r
166    }
167
168    /// Returns a libtest-mimic trial to run the test.
169    fn trial(
170        self,
171        resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
172    ) -> libtest_mimic::Trial {
173        let name = self.name();
174        let is_unstable = name.ends_with("_unstable");
175        libtest_mimic::Trial::test(name, move || {
176            map_test_result(self.run(resolve), &self.name(), is_unstable)
177        })
178    }
179}
180
181/// Maps a test result to a libtest-mimic result, suppressing failures for
182/// unstable tests.
183fn map_test_result(
184    result: anyhow::Result<()>,
185    test_name: &str,
186    is_unstable: bool,
187) -> Result<(), libtest_mimic::Failed> {
188    match result {
189        Ok(()) => Ok(()),
190        Err(err) if is_unstable => {
191            tracing::warn!("unstable test failed (non-gating): {test_name}: {err:#}");
192            Ok(())
193        }
194        Err(err) => Err(format!("{err:#}").into()),
195    }
196}
197
198/// A test that can be run.
199///
200/// Register it to be run with [`test!`] or [`multitest!`].
201pub trait RunTest: Send {
202    /// The type of artifacts required by the test.
203    type Artifacts;
204
205    /// The leaf name of the test.
206    ///
207    /// To produce the full test name, this will be prefixed with the module
208    /// name where the test is defined.
209    fn leaf_name(&self) -> &str;
210    /// Returns the artifacts required by the test.
211    ///
212    /// Returns `None` if this test makes no sense for this host environment
213    /// (e.g., an x86_64 test on an aarch64 host) and should be left out of the
214    /// test list.
215    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts>;
216    /// Runs the test, which has been assigned `name`, with the given
217    /// `artifacts`.
218    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()>;
219    /// Returns the host requirements of the current test, if any.
220    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
221}
222
223trait DynRunTest: Send {
224    fn leaf_name(&self) -> &str;
225    fn artifact_requirements(&self) -> Option<TestArtifactRequirements>;
226    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()>;
227    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
228}
229
230impl<T: RunTest> DynRunTest for T {
231    fn leaf_name(&self) -> &str {
232        self.leaf_name()
233    }
234
235    fn artifact_requirements(&self) -> Option<TestArtifactRequirements> {
236        let mut requirements = TestArtifactRequirements::new();
237        self.resolve(&ArtifactResolver::collector(&mut requirements))?;
238        Some(requirements)
239    }
240
241    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()> {
242        let artifacts = self
243            .resolve(&ArtifactResolver::resolver(artifacts))
244            .context("test should have been skipped")?;
245        self.run(params, artifacts)
246    }
247
248    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
249        self.host_requirements()
250    }
251}
252
253/// Parameters passed to a [`RunTest`] when it is run.
254pub struct PetriTestParams<'a> {
255    /// The name of the running test.
256    pub test_name: &'a str,
257    /// The logger for the test.
258    pub logger: &'a PetriLogSource,
259    /// Any hooks that want to run after the test completes.
260    pub post_test_hooks: &'a mut Vec<PetriPostTestHook>,
261}
262
263/// A post-test hook to be run after the test completes, regardless of if it
264/// succeeds or fails.
265pub struct PetriPostTestHook {
266    /// The name of the hook.
267    name: String,
268    /// The hook function.
269    hook: Box<dyn FnOnce(bool) -> anyhow::Result<()>>,
270}
271
272impl PetriPostTestHook {
273    pub fn new(name: String, hook: impl FnOnce(bool) -> anyhow::Result<()> + 'static) -> Self {
274        Self {
275            name,
276            hook: Box::new(hook),
277        }
278    }
279
280    pub fn name(&self) -> &str {
281        &self.name
282    }
283
284    pub fn run(self, test_passed: bool) -> anyhow::Result<()> {
285        (self.hook)(test_passed)
286    }
287}
288
289/// A test defined by an artifact resolver function and a run function.
290pub struct SimpleTest<A, F> {
291    leaf_name: &'static str,
292    resolve: A,
293    run: F,
294    /// Optional test requirements
295    pub host_requirements: Option<TestCaseRequirements>,
296}
297
298impl<A, AR, F, E> SimpleTest<A, F>
299where
300    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
301    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
302    E: Into<anyhow::Error>,
303{
304    /// Returns a new test with the given `leaf_name`, `resolve`, `run` functions,
305    /// and optional requirements.
306    pub fn new(
307        leaf_name: &'static str,
308        resolve: A,
309        run: F,
310        host_requirements: Option<TestCaseRequirements>,
311    ) -> Self {
312        SimpleTest {
313            leaf_name,
314            resolve,
315            run,
316            host_requirements,
317        }
318    }
319}
320
321impl<A, AR, F, E> RunTest for SimpleTest<A, F>
322where
323    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
324    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
325    E: Into<anyhow::Error>,
326{
327    type Artifacts = AR;
328
329    fn leaf_name(&self) -> &str {
330        self.leaf_name
331    }
332
333    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts> {
334        (self.resolve)(resolver)
335    }
336
337    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()> {
338        (self.run)(params, artifacts).map_err(Into::into)
339    }
340
341    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
342        self.host_requirements.as_ref()
343    }
344}
345
346#[derive(clap::Parser)]
347struct Options {
348    /// Lists the required artifacts for all tests.
349    #[clap(long)]
350    list_required_artifacts: bool,
351    #[clap(flatten)]
352    inner: libtest_mimic::Arguments,
353}
354
355/// Entry point for test binaries.
356pub fn test_main(
357    resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
358) -> ! {
359    let mut args = <Options as clap::Parser>::parse();
360    if args.list_required_artifacts {
361        // FUTURE: write this in a machine readable format.
362        for test in Test::all() {
363            println!("{}:", test.name());
364            for artifact in test.artifact_requirements.required_artifacts() {
365                println!("required: {artifact:?}");
366            }
367            for artifact in test.artifact_requirements.optional_artifacts() {
368                println!("optional: {artifact:?}");
369            }
370            println!();
371        }
372        std::process::exit(0);
373    }
374
375    // Always just use one thread to avoid interleaving logs and to avoid using
376    // too many resources. These tests are usually run under nextest, which will
377    // run them in parallel in separate processes with appropriate concurrency
378    // limits.
379    if !matches!(args.inner.test_threads, None | Some(1)) {
380        eprintln!("warning: ignoring value passed to --test-threads, using 1");
381    }
382    args.inner.test_threads = Some(1);
383
384    // Create the host context once to avoid repeated expensive queries
385    let host_context = futures::executor::block_on(HostContext::new());
386
387    let trials = Test::all()
388        .map(|test| {
389            let can_run = can_run_test_with_context(test.test.0.host_requirements(), &host_context);
390            test.trial(resolve).with_ignored_flag(!can_run)
391        })
392        .collect();
393
394    libtest_mimic::run(&args.inner, trials).exit();
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400
401    #[test]
402    fn stable_test_failure_propagates() {
403        let result: anyhow::Result<()> = Err(anyhow::anyhow!("boom"));
404        assert!(map_test_result(result, "some_test", false).is_err());
405    }
406
407    #[test]
408    fn unstable_test_failure_is_suppressed() {
409        let result: anyhow::Result<()> = Err(anyhow::anyhow!("boom"));
410        assert!(map_test_result(result, "some_test_unstable", true).is_ok());
411    }
412
413    #[test]
414    fn unstable_test_success_passes_through() {
415        let result: anyhow::Result<()> = Ok(());
416        assert!(map_test_result(result, "some_test_unstable", true).is_ok());
417    }
418
419    #[test]
420    fn mid_name_unstable_is_not_suppressed() {
421        // _unstable appearing mid-name should not be treated as unstable.
422        // The caller (trial) uses ends_with("_unstable") to set is_unstable,
423        // so a mid-name occurrence means is_unstable=false.
424        let result: anyhow::Result<()> = Err(anyhow::anyhow!("boom"));
425        assert!(map_test_result(result, "foo_unstable_bar", false).is_err());
426    }
427}