Skip to main content

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 futures::FutureExt as _;
33use pal_async::DefaultDriver;
34use pal_async::DefaultPool;
35use petri_artifacts_core::ArtifactResolver;
36use petri_artifacts_core::RemoteAccess;
37use std::collections::BTreeMap;
38use std::panic::AssertUnwindSafe;
39use std::panic::catch_unwind;
40use test_macro_support::TESTS;
41
42/// Defines a single test named after the async run function `$f`, using `$req`
43/// to resolve its artifacts. See [`SimpleTest::new_async`].
44#[macro_export]
45macro_rules! test {
46    ($f:ident, $req:expr) => {
47        $crate::multitest!(vec![
48            $crate::SimpleTest::new_async(stringify!($f), $req, $f).into()
49        ]);
50    };
51}
52
53/// Defines a single unstable test named after the async run function `$f`,
54/// using `$req` to resolve its artifacts.
55///
56/// `$reason` documents why the test is unstable and is logged when an unstable
57/// failure is ignored.
58#[macro_export]
59macro_rules! unstable_test {
60    ($f:ident, $req:expr, $reason:expr) => {
61        $crate::multitest!(vec![
62            $crate::SimpleTest::new_async(stringify!($f), $req, $f)
63                .unstable($reason)
64                .into()
65        ]);
66    };
67}
68
69/// Defines a single test named after the synchronous run function `$f`, using
70/// `$req` to resolve its artifacts. See [`SimpleTest::new_sync`].
71#[macro_export]
72macro_rules! test_sync {
73    ($f:ident, $req:expr) => {
74        $crate::multitest!(vec![
75            $crate::SimpleTest::new_sync(stringify!($f), $req, $f).into()
76        ]);
77    };
78}
79
80/// Defines a single unstable test named after the synchronous run function
81/// `$f`, using `$req` to resolve its artifacts.
82///
83/// `$reason` documents why the test is unstable and is logged when an unstable
84/// failure is ignored.
85#[macro_export]
86macro_rules! unstable_test_sync {
87    ($f:ident, $req:expr, $reason:expr) => {
88        $crate::multitest!(vec![
89            $crate::SimpleTest::new_sync(stringify!($f), $req, $f)
90                .unstable($reason)
91                .into()
92        ]);
93    };
94}
95
96/// Defines a set of tests from an expression evaluating to a
97/// [`Vec<TestCase>`](TestCase).
98#[macro_export]
99macro_rules! multitest {
100    ($tests:expr) => {
101        const _: () = {
102            use $crate::test_macro_support::linkme;
103            #[linkme::distributed_slice($crate::test_macro_support::TESTS)]
104            #[linkme(crate = linkme)]
105            static TEST: Option<fn() -> (&'static str, Vec<$crate::TestCase>)> =
106                Some(|| (module_path!(), $tests));
107        };
108    };
109}
110
111/// A single test case.
112pub struct TestCase(Box<dyn DynRunTest>);
113
114impl TestCase {
115    /// Creates a new test case from a value that implements [`RunTest`].
116    pub fn new(test: impl 'static + RunTest) -> Self {
117        Self(Box::new(test))
118    }
119}
120
121impl<T: 'static + RunTest> From<T> for TestCase {
122    fn from(test: T) -> Self {
123        Self::new(test)
124    }
125}
126
127/// A single test, with module name.
128struct Test {
129    module: &'static str,
130    test: TestCase,
131    artifact_requirements: TestArtifactRequirements,
132}
133
134impl Test {
135    /// Returns all the tests defined in this crate.
136    fn all() -> impl Iterator<Item = Self> {
137        TESTS.iter().flatten().flat_map(|f| {
138            let (module, tests) = f();
139            tests.into_iter().filter_map(move |test| {
140                let mut artifact_requirements = test.0.artifact_requirements()?;
141                // All tests require the log directory.
142                artifact_requirements.require(
143                    petri_artifacts_common::artifacts::TEST_LOG_DIRECTORY,
144                    RemoteAccess::LocalOnly,
145                    false,
146                    None,
147                );
148                Some(Self {
149                    module,
150                    artifact_requirements,
151                    test,
152                })
153            })
154        })
155    }
156
157    /// Returns the name of the test.
158    fn name(&self) -> String {
159        // Strip the crate name from the module path, for consistency with libtest.
160        match self.module.split_once("::") {
161            Some((_crate_name, rest)) => format!("{}::{}", rest, self.test.0.leaf_name()),
162            None => self.test.0.leaf_name().to_owned(),
163        }
164    }
165
166    fn run(
167        &self,
168        resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
169    ) -> anyhow::Result<()> {
170        let name = self.name();
171        let artifacts = resolve(&name, self.artifact_requirements.clone())
172            .context("failed to resolve artifacts")?;
173        let output_dir = artifacts.get(petri_artifacts_common::artifacts::TEST_LOG_DIRECTORY);
174        let is_nested_host = self
175            .artifact_requirements
176            .nested_artifacts()
177            .next()
178            .is_some();
179        let logger = try_init_tracing(
180            output_dir,
181            tracing::level_filters::LevelFilter::DEBUG,
182            is_nested_host.then(|| "l1_".into()),
183        )
184        .context("failed to initialize tracing")?;
185        // Record the test's identity up front, so that a test which is killed
186        // or crashes before reporting a result is still identifiable.
187        logger.log_test_start(&name);
188        let mut post_test_hooks = Vec::new();
189
190        // A process that is faulted or killed writes nothing to its own logs,
191        // so the host's error reporting events are the only record that it
192        // crashed and the only source of the Watson report ID needed to find
193        // the dump.
194        #[cfg(windows)]
195        post_test_hooks.push(collect_watson_events_hook(logger.clone()));
196
197        // Catch test panics in order to cleanly log the panic result. Without
198        // this, `libtest_mimic` will report the panic to stdout and fail the
199        // test, but the details won't end up in our per-test JSON log.
200        let r = catch_unwind(AssertUnwindSafe(|| {
201            self.test.0.run(
202                PetriTestParams {
203                    test_name: &name,
204                    log_source: &logger,
205                    post_test_hooks: &mut post_test_hooks,
206                },
207                &artifacts,
208            )
209        }));
210        let r = r.unwrap_or_else(|err| {
211            // The error from `catch_unwind` is almost always either a
212            // `&str` or a `String`, since that's what `panic!` produces.
213            let msg = err
214                .downcast_ref::<&str>()
215                .copied()
216                .or_else(|| err.downcast_ref::<String>().map(|x| x.as_str()));
217
218            let err = if let Some(msg) = msg {
219                anyhow::anyhow!("test panicked: {msg}")
220            } else {
221                anyhow::anyhow!("test panicked (unknown payload type)")
222            };
223            Err(err)
224        });
225        logger.log_test_result(&r, self.test.0.unstable().is_some());
226
227        for hook in post_test_hooks {
228            tracing::info!(name = hook.name(), "Running post-test hook");
229            if let Err(e) = hook.run(r.is_ok()) {
230                tracing::error!(
231                    error = e.as_ref() as &dyn std::error::Error,
232                    "Post-test hook failed"
233                );
234            } else {
235                tracing::info!("Post-test hook completed successfully");
236            }
237        }
238
239        r
240    }
241
242    /// Returns a libtest-mimic trial to run the test.
243    fn trial(
244        self,
245        resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
246    ) -> libtest_mimic::Trial {
247        libtest_mimic::Trial::test(self.name(), move || {
248            let unstable = self.test.0.unstable();
249            match self.run(resolve) {
250                Ok(()) => Ok(()),
251                Err(err) => {
252                    let Some(reason) = unstable else {
253                        return Err(format!("{err:#}").into());
254                    };
255                    if std::env::var(petri_artifacts_core::env::PETRI_IGNORE_UNSTABLE_FAILURES)
256                        .ok()
257                        .is_some_and(|v| !v.is_empty() && v != "0")
258                    {
259                        tracing::warn!(reason, "ignoring unstable test failure: {err:#}");
260                        return Ok(());
261                    }
262                    Err(format!("unstable test failed (reason: {reason}): {err:#}").into())
263                }
264            }
265        })
266    }
267}
268
269/// A test that can be run.
270///
271/// Register it to be run with [`test!`] or [`multitest!`].
272pub trait RunTest: Send {
273    /// The type of artifacts required by the test.
274    type Artifacts;
275
276    /// The leaf name of the test.
277    ///
278    /// To produce the full test name, this will be prefixed with the module
279    /// name where the test is defined.
280    fn leaf_name(&self) -> &str;
281    /// Returns the artifacts required by the test.
282    ///
283    /// Returns `None` if this test makes no sense for this host environment
284    /// (e.g., an x86_64 test on an aarch64 host) and should be left out of the
285    /// test list.
286    fn resolve(&self, resolver: ArtifactResolver<'_>) -> Option<Self::Artifacts>;
287    /// Runs the test, which has been assigned `name`, with the given
288    /// `artifacts`.
289    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()>;
290    /// Returns the host requirements of the current test, if any.
291    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
292    /// If this test is unstable, the reason why; `None` if stable.
293    fn unstable(&self) -> Option<&str>;
294    /// Whether this test is ignored (skipped by default, like a libtest
295    /// `#[ignore]` test).
296    fn ignored(&self) -> bool;
297}
298
299trait DynRunTest: Send {
300    fn leaf_name(&self) -> &str;
301    fn artifact_requirements(&self) -> Option<TestArtifactRequirements>;
302    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()>;
303    fn host_requirements(&self) -> Option<&TestCaseRequirements>;
304    fn unstable(&self) -> Option<&str>;
305    fn ignored(&self) -> bool;
306}
307
308impl<T: RunTest> DynRunTest for T {
309    fn leaf_name(&self) -> &str {
310        self.leaf_name()
311    }
312
313    fn artifact_requirements(&self) -> Option<TestArtifactRequirements> {
314        let mut requirements = TestArtifactRequirements::new();
315        self.resolve(ArtifactResolver::collector(&mut requirements))?;
316        Some(requirements)
317    }
318
319    fn run(&self, params: PetriTestParams<'_>, artifacts: &TestArtifacts) -> anyhow::Result<()> {
320        let artifacts = self
321            .resolve(ArtifactResolver::resolver(artifacts))
322            .context("test should have been skipped")?;
323        self.run(params, artifacts)
324    }
325
326    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
327        self.host_requirements()
328    }
329
330    fn unstable(&self) -> Option<&str> {
331        self.unstable()
332    }
333
334    fn ignored(&self) -> bool {
335        self.ignored()
336    }
337}
338
339/// Parameters passed to a [`RunTest`] when it is run.
340pub struct PetriTestParams<'a> {
341    /// The name of the running test.
342    pub test_name: &'a str,
343    /// The logger for the test.
344    pub log_source: &'a PetriLogSource,
345    /// Any hooks that want to run after the test completes.
346    pub post_test_hooks: &'a mut Vec<PetriPostTestHook>,
347}
348
349/// A post-test hook to be run after the test completes, regardless of if it
350/// succeeds or fails.
351pub struct PetriPostTestHook {
352    /// The name of the hook.
353    name: String,
354    /// The hook function.
355    hook: Box<dyn FnOnce(bool) -> anyhow::Result<()>>,
356}
357
358impl PetriPostTestHook {
359    pub fn new(name: String, hook: impl FnOnce(bool) -> anyhow::Result<()> + 'static) -> Self {
360        Self {
361            name,
362            hook: Box::new(hook),
363        }
364    }
365
366    pub fn name(&self) -> &str {
367        &self.name
368    }
369
370    pub fn run(self, test_passed: bool) -> anyhow::Result<()> {
371        (self.hook)(test_passed)
372    }
373}
374
375/// Returns a hook that, if the test failed, writes the Windows Error Reporting
376/// and Azure Watson events from the test's execution window to
377/// `watson_events.log`.
378#[cfg(windows)]
379fn collect_watson_events_hook(logger: PetriLogSource) -> PetriPostTestHook {
380    let start_time = jiff::Timestamp::now();
381    PetriPostTestHook::new("collect watson events".into(), move |test_passed| {
382        if test_passed {
383            return Ok(());
384        }
385        let events =
386            futures::executor::block_on(crate::vm::hyperv::powershell::watson_events(&start_time));
387        if events.is_empty() {
388            return Ok(());
389        }
390        let log_file = logger.log_file("watson_events")?;
391        for event in events {
392            event.write_to(&log_file);
393        }
394        Ok(())
395    })
396}
397
398/// A test defined by an artifact resolver function and a run function.
399pub struct SimpleTest<A, F> {
400    leaf_name: &'static str,
401    resolve: A,
402    run: F,
403    /// Optional test requirements
404    pub host_requirements: Option<TestCaseRequirements>,
405    unstable: Option<&'static str>,
406    ignored: bool,
407    remote_policy: RemoteAccess,
408}
409
410impl<A, AR, F, E> SimpleTest<A, F>
411where
412    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
413    F: 'static + Send + AsyncFn(PetriTestParams<'_>, DefaultDriver, AR) -> Result<(), E>,
414    E: Into<anyhow::Error>,
415{
416    /// Returns a new test with the given `leaf_name`, `resolve`, and `run`
417    /// functions, whose `run` function is async, running it on a task
418    /// pool owned by petri.
419    ///
420    /// The test defaults to stable, not ignored, with no host requirements and
421    /// a [`RemoteAccess::LocalOnly`] policy. Use the builder methods
422    /// ([`requirements`](Self::requirements), [`unstable`](Self::unstable),
423    /// [`ignore`](Self::ignore), [`remote_access`](Self::remote_access)) to
424    /// override these.
425    pub fn new_async(
426        leaf_name: &'static str,
427        resolve: A,
428        run: F,
429    ) -> SimpleTest<A, impl 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>> {
430        SimpleTest::new_sync(leaf_name, resolve, move |params, artifacts| {
431            let mut pool = DefaultPool::named(std::thread::current().name().unwrap_or(leaf_name));
432            let driver = pool.driver();
433            // The inner catch keeps a panicking test body from unwinding through
434            // the pool; the outer one catches panics raised by spawned tasks,
435            // such as the VM's timeout watchdog.
436            let r = catch_unwind(AssertUnwindSafe(|| {
437                pool.run_until(
438                    AssertUnwindSafe(run(params, driver.clone(), artifacts)).catch_unwind(),
439                )
440            }));
441            // Let the diagnostic tasks the VM spawns as it is dropped finish
442            // before the failure is reported.
443            drop(driver);
444            pool.run();
445            match r.and_then(|r| r) {
446                Ok(r) => r,
447                Err(panic) => std::panic::resume_unwind(panic),
448            }
449        })
450    }
451}
452
453impl<A, AR, F, E> SimpleTest<A, F>
454where
455    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
456    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
457    E: Into<anyhow::Error>,
458{
459    /// Returns a new test with the given `leaf_name`, `resolve`, and `run`
460    /// functions.
461    ///
462    /// Use [`new_async`](Self::new_async) instead if the test constructs a
463    /// [`PetriVm`](crate::PetriVm); its diagnostics on failure rely on petri
464    /// owning the task pool.
465    ///
466    /// The test defaults to stable, not ignored, with no host requirements and
467    /// a [`RemoteAccess::LocalOnly`] policy. Use the builder methods
468    /// ([`requirements`](Self::requirements), [`unstable`](Self::unstable),
469    /// [`ignore`](Self::ignore), [`remote_access`](Self::remote_access)) to
470    /// override these.
471    pub fn new_sync(leaf_name: &'static str, resolve: A, run: F) -> Self {
472        SimpleTest {
473            leaf_name,
474            resolve,
475            run,
476            host_requirements: None,
477            unstable: None,
478            ignored: false,
479            remote_policy: RemoteAccess::LocalOnly,
480        }
481    }
482}
483
484impl<A, F> SimpleTest<A, F> {
485    /// Sets the host requirements that must be satisfied for this test to run.
486    pub fn requirements(mut self, requirements: TestCaseRequirements) -> Self {
487        self.host_requirements = Some(requirements);
488        self
489    }
490
491    /// Marks this test as unstable. When `PETRI_IGNORE_UNSTABLE_FAILURES` is
492    /// set (as it is in CI), a failure of this test is logged and ignored
493    /// rather than failing the run; otherwise it fails like any other test.
494    ///
495    /// `reason` documents why the test is unstable and is logged when an
496    /// unstable failure is ignored.
497    pub fn unstable(mut self, reason: &'static str) -> Self {
498        self.unstable = Some(reason);
499        self
500    }
501
502    /// Marks this test as ignored: it is skipped by default and only runs when
503    /// explicitly requested (like a libtest `#[ignore]` test).
504    pub fn ignore(mut self) -> Self {
505        self.ignored = true;
506        self
507    }
508
509    /// Sets the remote-access policy used when resolving artifacts.
510    pub fn remote_access(mut self, policy: RemoteAccess) -> Self {
511        self.remote_policy = policy;
512        self
513    }
514}
515
516impl<A, AR, F, E> RunTest for SimpleTest<A, F>
517where
518    A: 'static + Send + Fn(&ArtifactResolver<'_>) -> Option<AR>,
519    F: 'static + Send + Fn(PetriTestParams<'_>, AR) -> Result<(), E>,
520    E: Into<anyhow::Error>,
521{
522    type Artifacts = AR;
523
524    fn leaf_name(&self) -> &str {
525        self.leaf_name
526    }
527
528    fn resolve(&self, mut resolver: ArtifactResolver<'_>) -> Option<Self::Artifacts> {
529        resolver.set_remote_policy(self.remote_policy);
530        (self.resolve)(&resolver)
531    }
532
533    fn run(&self, params: PetriTestParams<'_>, artifacts: Self::Artifacts) -> anyhow::Result<()> {
534        (self.run)(params, artifacts).map_err(Into::into)
535    }
536
537    fn host_requirements(&self) -> Option<&TestCaseRequirements> {
538        self.host_requirements.as_ref()
539    }
540
541    fn unstable(&self) -> Option<&str> {
542        self.unstable
543    }
544
545    fn ignored(&self) -> bool {
546        self.ignored
547    }
548}
549
550#[derive(clap::Parser)]
551struct Options {
552    /// Lists the required artifacts for all tests in JSON format.
553    /// Use --tests-from-stdin to query artifacts for specific tests.
554    #[clap(long)]
555    list_required_artifacts: bool,
556    /// When used with --list-required-artifacts, read exact test names from
557    /// stdin (one per line) to query artifacts for specific tests only.
558    ///
559    /// Even though users can use nextest's filter logic to run a subset of
560    /// tests, due to nextest's architecture of running one test per binary we
561    /// cannot accept a nextest filter here directly. Instead, vmm-tests-run
562    /// must first call nextest with the desired filter to determine the exact
563    /// test names to pass via stdin, then ask petri what artifacts are required
564    /// for those tests.
565    #[clap(long, requires = "list_required_artifacts")]
566    tests_from_stdin: bool,
567    #[clap(flatten)]
568    inner: libtest_mimic::Arguments,
569}
570
571/// Entry point for test binaries.
572pub fn test_main(
573    resolve: fn(&str, TestArtifactRequirements) -> anyhow::Result<TestArtifacts>,
574) -> ! {
575    let mut args = <Options as clap::Parser>::parse();
576    if args.list_required_artifacts {
577        use std::collections::BTreeSet;
578
579        // Collect all artifacts from tests (all tests, or those specified via stdin)
580        let mut required_set = BTreeSet::new();
581        let mut optional_set = BTreeSet::new();
582        let mut nested_set: BTreeMap<String, Vec<String>> = BTreeMap::new();
583
584        // If reading test names from stdin, collect them into a set for exact matching
585        let stdin_tests: Option<BTreeSet<String>> = if args.tests_from_stdin {
586            use std::io::BufRead;
587            let stdin = std::io::stdin();
588            let tests: BTreeSet<String> = stdin
589                .lock()
590                .lines()
591                .map_while(Result::ok)
592                .filter(|line| !line.is_empty())
593                .collect();
594            if tests.is_empty() {
595                eprintln!("warning: no test names provided on stdin");
596            }
597            Some(tests)
598        } else {
599            None
600        };
601
602        for test in Test::all() {
603            let name = test.name();
604
605            // If reading from stdin, do exact matching; otherwise include all tests
606            let matches = match stdin_tests {
607                Some(ref stdin_tests) => stdin_tests.contains(&name),
608                None => true,
609            };
610
611            if matches {
612                for artifact in test.artifact_requirements.required_artifacts() {
613                    required_set.insert(artifact.global_unique_id().to_string());
614                }
615                for artifact in test.artifact_requirements.optional_artifacts() {
616                    optional_set.insert(artifact.global_unique_id().to_string());
617                }
618                for (artifact, test) in test.artifact_requirements.nested_artifacts() {
619                    nested_set
620                        .entry(artifact.global_unique_id().to_string())
621                        .or_default()
622                        .push(test.to_string());
623                }
624            }
625        }
626
627        // Remove from optional any artifacts that are required
628        let optional_set: BTreeSet<_> = optional_set.difference(&required_set).cloned().collect();
629
630        let output = petri_artifacts_core::ArtifactListOutput {
631            required: required_set.into_iter().collect(),
632            optional: optional_set.into_iter().collect(),
633            nested: nested_set.into_iter().collect(),
634        };
635
636        println!(
637            "{}",
638            serde_json::to_string(&output).expect("JSON serialization failed")
639        );
640        std::process::exit(0);
641    }
642
643    // Always just use one thread to avoid interleaving logs and to avoid using
644    // too many resources. These tests are usually run under nextest, which will
645    // run them in parallel in separate processes with appropriate concurrency
646    // limits.
647    if !matches!(args.inner.test_threads, None | Some(1)) {
648        eprintln!("warning: ignoring value passed to --test-threads, using 1");
649    }
650    args.inner.test_threads = Some(1);
651
652    // Create the host context once to avoid repeated expensive queries
653    let host_context = futures::executor::block_on(HostContext::new());
654
655    let trials = Test::all()
656        .map(|test| {
657            let can_run = can_run_test_with_context(test.test.0.host_requirements(), &host_context);
658            let ignored = test.test.0.ignored();
659            test.trial(resolve).with_ignored_flag(!can_run || ignored)
660        })
661        .collect();
662
663    libtest_mimic::run(&args.inner, trials).exit();
664}