Skip to main content

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}