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

petri/
test.rs
