#
cargo_test_support/
containers.rs

    1//! Support for testing using Docker containers.
    2//!
    3//! The [`Container`] type is a builder for configuring a container to run.
    4//! After you call `launch`, you can use the [`ContainerHandle`] to interact
    5//! with the running container.
    6//!
    7//! Tests using containers must use `#[cargo_test(container_test)]` to disable
    8//! them unless the `CARGO_CONTAINER_TESTS` environment variable is set.
    9
    10use cargo_util::ProcessBuilder;
    11use std::collections::HashMap;
    12use std::io::Read;
    13use std::path::PathBuf;
    14use std::process::Command;
    15use std::sync::Mutex;
    16use std::sync::atomic::{AtomicUsize, Ordering};
    17use tar::Header;
    18
    19/// A builder for configuring a container to run.
    20pub struct Container {
    21    /// The host directory that forms the basis of the Docker image.
    22    build_context: PathBuf,
    23    /// Files to copy over to the image.
    24    files: Vec<MkFile>,
    25}
    26
    27/// A handle to a running container.
    28///
    29/// You can use this to interact with the container.
    30pub struct ContainerHandle {
    31    /// The name of the container.
    32    name: String,
    33    /// Port mappings of `container_port` to `host_port` for ports exposed via EXPOSE.
    34    pub port_mappings: HashMap<u16, u16>,
    35}
    36
    37impl Container {
    38    pub fn new(context_dir: &str) -> Container {
    39        assert!(std::env::var_os("CARGO_CONTAINER_TESTS").is_some());
    40        let mut build_context = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    41        build_context.push("containers");
    42        build_context.push(context_dir);
    43        Container {
    44            build_context,
    45            files: Vec::new(),
    46        }
    47    }
    48
    49    /// Adds a file to be copied into the container.
    50    pub fn file(mut self, file: MkFile) -> Self {
    51        self.files.push(file);
    52        self
    53    }
    54
    55    /// Starts the container.
    56    pub fn launch(mut self) -> ContainerHandle {
    57        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
    58
    59        let id = NEXT_ID.fetch_add(1, Ordering::SeqCst);
    60        let name = format!("cargo_test_{id}");
    61        remove_if_exists(&name);
    62        self.create_container(&name);
    63        self.copy_files(&name);
    64        self.start_container(&name);
    65        let info = self.container_inspect(&name);
    66        let port_mappings = self.port_mappings(&info);
    67        self.wait_till_ready(&port_mappings);
    68
    69        ContainerHandle {
    70            name,
    71            port_mappings,
    72        }
    73    }
    74
    75    fn create_container(&self, name: &str) {
    76        static BUILD_LOCK: Mutex<()> = Mutex::new(());
    77
    78        let image_base = self.build_context.file_name().unwrap();
    79        let image_name = format!("cargo-test-{}", image_base.to_str().unwrap());
    80        let _lock = BUILD_LOCK
    81            .lock()
    82            .map_err(|_| panic!("previous docker build failed, unable to run test"));
    83        ProcessBuilder::new("docker")
    84            .args(&["build", "--tag", image_name.as_str()])
    85            .arg(&self.build_context)
    86            .exec_with_output()
    87            .unwrap();
    88
    89        ProcessBuilder::new("docker")
    90            .args(&[
    91                "container",
    92                "create",
    93                "--publish-all",
    94                "--rm",
    95                "--name",
    96                name,
    97            ])
    98            .arg(image_name)
    99            .exec_with_output()
    100            .unwrap();
    101    }
    102
    103    fn copy_files(&mut self, name: &str) {
    104        if self.files.is_empty() {
    105            return;
    106        }
    107        let mut ar = tar::Builder::new(Vec::new());
    108        ar.sparse(false);
    109        let files = std::mem::replace(&mut self.files, Vec::new());
    110        for mut file in files {
    111            ar.append_data(&mut file.header, &file.path, file.contents.as_slice())
    112                .unwrap();
    113        }
    114        let ar = ar.into_inner().unwrap();
    115        ProcessBuilder::new("docker")
    116            .args(&["cp", "-"])
    117            .arg(format!("{name}:/"))
    118            .stdin(ar)
    119            .exec_with_output()
    120            .unwrap();
    121    }
    122
    123    fn start_container(&self, name: &str) {
    124        ProcessBuilder::new("docker")
    125            .args(&["container", "start"])
    126            .arg(name)
    127            .exec_with_output()
    128            .unwrap();
    129    }
    130
    131    fn container_inspect(&self, name: &str) -> serde_json::Value {
    132        let output = ProcessBuilder::new("docker")
    133            .args(&["inspect", name])
    134            .exec_with_output()
    135            .unwrap();
    136        serde_json::from_slice(&output.stdout).unwrap()
    137    }
    138
    139    /// Returns the mapping of container_port->host_port for ports that were
    140    /// exposed with EXPOSE.
    141    fn port_mappings(&self, info: &serde_json::Value) -> HashMap<u16, u16> {
    142        info[0]["NetworkSettings"]["Ports"]
    143            .as_object()
    144            .unwrap()
    145            .iter()
    146            .map(|(key, value)| {
    147                let key = key
    148                    .strip_suffix("/tcp")
    149                    .expect("expected TCP only ports")
    150                    .parse()
    151                    .unwrap();
    152                let values = value.as_array().unwrap();
    153                let value = values
    154                    .iter()
    155                    .find(|value| value["HostIp"].as_str().unwrap() == "0.0.0.0")
    156                    .expect("expected localhost IP");
    157                let host_port = value["HostPort"].as_str().unwrap().parse().unwrap();
    158                (key, host_port)
    159            })
    160            .collect()
    161    }
    162
    163    fn wait_till_ready(&self, port_mappings: &HashMap<u16, u16>) {
    164        for port in port_mappings.values() {
    165            let mut ok = false;
    166            for _ in 0..30 {
    167                match std::net::TcpStream::connect(format!("127.0.0.1:{port}")) {
    168                    Ok(_) => {
    169                        ok = true;
    170                        break;
    171                    }
    172                    Err(e) => {
    173                        if e.kind() != std::io::ErrorKind::ConnectionRefused {
    174                            panic!("unexpected localhost connection error: {e:?}");
    175                        }
    176                        std::thread::sleep(std::time::Duration::new(1, 0));
    177                    }
    178                }
    179            }
    180            if !ok {
    181                panic!("no listener on localhost port {port}");
    182            }
    183        }
    184    }
    185}
    186
    187impl ContainerHandle {
    188    /// Executes a program inside a running container.
    189    pub fn exec(&self, args: &[&str]) -> std::process::Output {
    190        ProcessBuilder::new("docker")
    191            .args(&["container", "exec", &self.name])
    192            .args(args)
    193            .exec_with_output()
    194            .unwrap()
    195    }
    196
    197    /// Returns the contents of a file inside the container.
    198    pub fn read_file(&self, path: &str) -> String {
    199        let output = ProcessBuilder::new("docker")
    200            .args(&["cp", &format!("{}:{}", self.name, path), "-"])
    201            .exec_with_output()
    202            .unwrap();
    203        let mut ar = tar::Archive::new(output.stdout.as_slice());
    204        let mut entry = ar.entries().unwrap().next().unwrap().unwrap();
    205        let mut contents = String::new();
    206        entry.read_to_string(&mut contents).unwrap();
    207        contents
    208    }
    209}
    210
    211impl Drop for ContainerHandle {
    212    fn drop(&mut self) {
    213        // To help with debugging, this will keep the container alive.
    214        if std::env::var_os("CARGO_CONTAINER_TEST_KEEP").is_some() {
    215            return;
    216        }
    217        remove_if_exists(&self.name);
    218    }
    219}
    220
    221fn remove_if_exists(name: &str) {
    222    if let Err(e) = Command::new("docker")
    223        .args(&["container", "rm", "--force", name])
    224        .output()
    225    {
    226        panic!("failed to run docker: {e}");
    227    }
    228}
    229
    230/// Builder for configuring a file to copy into a container.
    231pub struct MkFile {
    232    path: String,
    233    contents: Vec<u8>,
    234    header: Header,
    235}
    236
    237impl MkFile {
    238    /// Defines a file to add to the container.
    239    ///
    240    /// This should be passed to `Container::file`.
    241    ///
    242    /// The path is the path inside the container to create the file.
    243    pub fn path(path: &str) -> MkFile {
    244        MkFile {
    245            path: path.to_string(),
    246            contents: Vec::new(),
    247            header: Header::new_gnu(),
    248        }
    249    }
    250
    251    pub fn contents(mut self, contents: impl Into<Vec<u8>>) -> Self {
    252        self.contents = contents.into();
    253        self.header.set_size(self.contents.len() as u64);
    254        self
    255    }
    256
    257    pub fn mode(mut self, mode: u32) -> Self {
    258        self.header.set_mode(mode);
    259        self
    260    }
    261
    262    pub fn uid(mut self, uid: u64) -> Self {
    263        self.header.set_uid(uid);
    264        self
    265    }
    266
    267    pub fn gid(mut self, gid: u64) -> Self {
    268        self.header.set_gid(gid);
    269        self
    270    }
    271}