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