hyperactor/test_utils/
process_assertion.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::future::Future;
10
11use nix::sys::wait::WaitStatus;
12use nix::sys::wait::waitpid;
13use nix::unistd::ForkResult;
14use nix::unistd::fork;
15
16/// Fork a child process, execute the given function in that process, and verify
17/// that the process exits with the given exit code.
18pub async fn assert_termination<F, Fut>(f: F, expected_code: i32) -> anyhow::Result<()>
19where
20    F: FnOnce() -> Fut,
21    Fut: Future<Output = ()>,
22{
23    // SAFETY: for unit test process assertion.
24    unsafe {
25        match fork() {
26            Ok(ForkResult::Parent { child, .. }) => match waitpid(child, None)? {
27                WaitStatus::Exited(_, exit_code) => {
28                    anyhow::ensure!(exit_code == expected_code);
29                    Ok(())
30                }
31                status => Err(anyhow::anyhow!(
32                    "didn't receive expected status. got: {:?}",
33                    status
34                )),
35            },
36            Ok(ForkResult::Child) => Ok(f().await),
37            Err(_) => Err(anyhow::anyhow!("fork failed")),
38        }
39    }
40}