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        libtest_mimic::Trial::test(self.name(), move || {
174            self.run(resolve).map_err(|err| format!("{err:#}").into())
175        })
176    }
177}
178
179/// A test that can be run.
180///
181/// Register it to be run with [`test!`] or [`multitest!`].
182pub trait RunTest: Send {
183    /// The type of artifacts required by the test.
184    type Artifacts;
185
186    /// The leaf name of the test.
187    ///
188    /// To produce the full test name, this will be prefixed with the module
189    /// name where the test is defined.
190    fn leaf_name(&self) -> &str;
191    /// Returns the artifacts required by the test.
192    ///
193    /// Returns `None` if this test makes no sense for this host environment
194    /// (e.g., an x86_64 test on an aarch64 host) and should be left out of the
195    /// test list.
196    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts>;
197    /// Runs the test, which has been assigned `name`, with the given
198    /// `artifacts`.
199    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()>;
200    /// Returns the host requirements of the current test, if any.
201    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
202}
203
204trait DynRunTest: Send {
205    fn leaf_name(&self) -> &str;
206    fn artifact_requirements(&self) -> Option<TestArtifactRequirements>;
207    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()>;
208    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
209}
210
211impl<T: RunTest> DynRunTest for T {
212    fn leaf_name(&self) -> &str {
213        self.leaf_name()
214    }
215
216    fn artifact_requirements(&self) -> Option<TestArtifactRequirements> {
217        let mut requirements = TestArtifactRequirements::new();
218        self.resolve(&ArtifactResolver::collector(&mut requirements))?;
219        Some(requirements)
220    }
221
222    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()> {
223        let artifacts = self
224            .resolve(&ArtifactResolver::resolver(artifacts))
225            .context("test should have been skipped")?;
226        self.run(params, artifacts)
227    }
228
229    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
230        self.host_requirements()
231    }
232}
233
234/// Parameters passed to a [`RunTest`] when it is run.
235pub struct PetriTestParams<'a> {
236    /// The name of the running test.
237    pub test_name: &'a str,
238    /// The logger for the test.
239    pub logger: &'a PetriLogSource,
240    /// Any hooks that want to run after the test completes.
241    pub post_test_hooks: &'a mut Vec<PetriPostTestHook>,
242}
243
244/// A post-test hook to be run after the test completes, regardless of if it
245/// succeeds or fails.
246pub struct PetriPostTestHook {
247    /// The name of the hook.
248    name: String,
249    /// The hook function.
250    hook: Box<dyn FnOnce(bool) -> anyhow::Result<()>>,
251}
252
253impl PetriPostTestHook {
254    pub fn new(name: String, hook: impl FnOnce(bool) -> anyhow::Result<()> + 'static) -> Self {
255        Self {
256            name,
257            hook: Box::new(hook),
258        }
259    }
260
261    pub fn name(&self) -> &str {
262        &self.name
263    }
264
265    pub fn run(self, test_passed: bool) -> anyhow::Result<()> {
266        (self.hook)(test_passed)
267    }
268}
269
270/// A test defined by an artifact resolver function and a run function.
271pub struct SimpleTest<A, F> {
272    leaf_name: &'static str,
273    resolve: A,
274    run: F,
275    /// Optional test requirements
276    pub host_requirements: Option<TestCaseRequirements>,
277}
278
279impl<A, AR, F, E> SimpleTest<A, F>
280where
281    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
282    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
283    E: Into<anyhow::Error>,
284{
285    /// Returns a new test with the given `leaf_name`, `resolve`, `run` functions,
286    /// and optional requirements.
287    pub fn new(
288        leaf_name: &'static str,
289        resolve: A,
290        run: F,
291        host_requirements: Option<TestCaseRequirements>,
292    ) -> Self {
293        SimpleTest {
294            leaf_name,
295            resolve,
296            run,
297            host_requirements,
298        }
299    }
300}
301
302impl<A, AR, F, E> RunTest for SimpleTest<A, F>
303where
304    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
305    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
306    E: Into<anyhow::Error>,
307{
308    type Artifacts = AR;
309
310    fn leaf_name(&self) -> &str {
311        self.leaf_name
312    }
313
314    fn resolve(&self, resolver: &ArtifactResolver<'_>) -> Option<Self::Artifacts> {
315        (self.resolve)(resolver)
316    }
317
318    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()> {
319        (self.run)(params, artifacts).map_err(Into::into)
320    }
321
322    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
323        self.host_requirements.as_ref()
324    }
325}
326
327#[derive(clap::Parser)]
328struct Options {
329    /// Lists the required artifacts for all tests.
330    #[clap(long)]
331    list_required_artifacts: bool,
332    #[clap(flatten)]
333    inner: libtest_mimic::Arguments,
334}
335
336/// Entry point for test binaries.
337pub fn test_main(
338    resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
339) -> ! {
340    let mut args = <Options as clap::Parser>::parse();
341    if args.list_required_artifacts {
342        // FUTURE: write this in a machine readable format.
343        for test in Test::all() {
344            println!("{}:", test.name());
345            for artifact in test.artifact_requirements.required_artifacts() {
346                println!("required: {artifact:?}");
347            }
348            for artifact in test.artifact_requirements.optional_artifacts() {
349                println!("optional: {artifact:?}");
350            }
351            println!();
352        }
353        std::process::exit(0);
354    }
355
356    // Always just use one thread to avoid interleaving logs and to avoid using
357    // too many resources. These tests are usually run under nextest, which will
358    // run them in parallel in separate processes with appropriate concurrency
359    // limits.
360    if !matches!(args.inner.test_threads, None | Some(1)) {
361        eprintln!("warning: ignoring value passed to --test-threads, using 1");
362    }
363    args.inner.test_threads = Some(1);
364
365    // Create the host context once to avoid repeated expensive queries
366    let host_context = futures::executor::block_on(HostContext::new());
367
368    let trials = Test::all()
369        .map(|test| {
370            let can_run = can_run_test_with_context(test.test.0.host_requirements(), &host_context);
371            test.trial(resolve).with_ignored_flag(!can_run)
372        })
373        .collect();
374
375    libtest_mimic::run(&args.inner, trials).exit();
376}
petri/test.rs

petri/
test.rs
