testcontainers/core/image/
image_ext.rs

1use std::time::Duration;
2
3#[cfg(feature = "device-requests")]
4use bollard::models::DeviceRequest;
5use bollard::models::ResourcesUlimits;
6
7use crate::{
8    core::{
9        copy::{CopyDataSource, CopyToContainer},
10        healthcheck::Healthcheck,
11        logs::consumer::LogConsumer,
12        CgroupnsMode, ContainerPort, Host, Mount, PortMapping, WaitFor,
13    },
14    ContainerRequest, Image,
15};
16
17#[cfg(feature = "reusable-containers")]
18#[derive(Eq, Copy, Clone, Debug, Default, PartialEq)]
19pub enum ReuseDirective {
20    #[default]
21    Never,
22    Always,
23    CurrentSession,
24}
25
26#[cfg(feature = "reusable-containers")]
27impl std::fmt::Display for ReuseDirective {
28    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        formatter.write_str(match self {
30            Self::Never => "never",
31            Self::Always => "always",
32            Self::CurrentSession => "current-session",
33        })
34    }
35}
36
37/// Represents an extension for the [`Image`] trait.
38/// Allows to override image defaults and container configuration.
39pub trait ImageExt<I: Image> {
40    /// Returns a new [`ContainerRequest`] with the specified (overridden) `CMD` ([`Image::cmd`]).
41    ///
42    /// # Examples
43    /// ```rust,no_run
44    /// use testcontainers::{GenericImage, ImageExt};
45    ///
46    /// let image = GenericImage::new("image", "tag");
47    /// let cmd = ["arg1", "arg2"];
48    /// let overridden_cmd = image.clone().with_cmd(cmd);
49    ///
50    /// assert!(overridden_cmd.cmd().eq(cmd));
51    ///
52    /// let another_container_req = image.with_cmd(cmd);
53    ///
54    /// assert!(another_container_req.cmd().eq(overridden_cmd.cmd()));
55    /// ```
56    fn with_cmd(self, cmd: impl IntoIterator<Item = impl Into<String>>) -> ContainerRequest<I>;
57
58    /// Overrides the fully qualified image name (consists of `{domain}/{owner}/{image}`).
59    /// Can be used to specify a custom registry or owner.
60    fn with_name(self, name: impl Into<String>) -> ContainerRequest<I>;
61
62    /// Overrides the image tag.
63    ///
64    /// There is no guarantee that the specified tag for an image would result in a
65    /// running container. Users of this API are advised to use this at their own risk.
66    fn with_tag(self, tag: impl Into<String>) -> ContainerRequest<I>;
67
68    /// Sets the container name.
69    fn with_container_name(self, name: impl Into<String>) -> ContainerRequest<I>;
70
71    /// Sets the platform the container will be run on.
72    ///
73    /// Platform in the format `os[/arch[/variant]]` used for image lookup.
74    ///
75    /// # Examples
76    ///
77    /// ```rust,no_run
78    /// use testcontainers::{GenericImage, ImageExt};
79    ///
80    /// let image = GenericImage::new("image", "tag")
81    ///     .with_platform("linux/amd64");
82    /// ```
83    fn with_platform(self, platform: impl Into<String>) -> ContainerRequest<I>;
84
85    /// Sets the network the container will be connected to.
86    fn with_network(self, network: impl Into<String>) -> ContainerRequest<I>;
87
88    /// Adds the specified label to the container.
89    ///
90    /// **Note**: all keys in the `org.testcontainers.*` namespace should be regarded
91    /// as reserved by `testcontainers` internally, and should not be expected or relied
92    /// upon to be applied correctly if supplied as a value for `key`.
93    fn with_label(self, key: impl Into<String>, value: impl Into<String>) -> ContainerRequest<I>;
94
95    /// Adds the specified labels to the container.
96    ///
97    /// **Note**: all keys in the `org.testcontainers.*` namespace should be regarded
98    /// as reserved by `testcontainers` internally, and should not be expected or relied
99    /// upon to be applied correctly if they are included in `labels`.
100    fn with_labels(
101        self,
102        labels: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
103    ) -> ContainerRequest<I>;
104
105    /// Adds an environment variable to the container.
106    fn with_env_var(self, name: impl Into<String>, value: impl Into<String>)
107        -> ContainerRequest<I>;
108
109    /// Adds a host to the container.
110    fn with_host(self, key: impl Into<String>, value: impl Into<Host>) -> ContainerRequest<I>;
111
112    /// Configures hostname for the container.
113    fn with_hostname(self, hostname: impl Into<String>) -> ContainerRequest<I>;
114
115    /// Adds a mount to the container.
116    fn with_mount(self, mount: impl Into<Mount>) -> ContainerRequest<I>;
117
118    /// Copies some source into the container as file
119    fn with_copy_to(
120        self,
121        target: impl Into<String>,
122        source: impl Into<CopyDataSource>,
123    ) -> ContainerRequest<I>;
124
125    /// Adds a port mapping to the container, mapping the host port to the container's internal port.
126    ///
127    /// # Examples
128    /// ```rust,no_run
129    /// use testcontainers::{GenericImage, ImageExt};
130    /// use testcontainers::core::IntoContainerPort;
131    ///
132    /// let image = GenericImage::new("image", "tag").with_mapped_port(8080, 80.tcp());
133    /// ```
134    fn with_mapped_port(self, host_port: u16, container_port: ContainerPort)
135        -> ContainerRequest<I>;
136
137    /// Declares a host port that should be reachable from inside the container.
138    #[cfg(feature = "host-port-exposure")]
139    fn with_exposed_host_port(self, port: u16) -> ContainerRequest<I>;
140
141    /// Declares multiple host ports that should be reachable from inside the container.
142    #[cfg(feature = "host-port-exposure")]
143    fn with_exposed_host_ports(self, ports: impl IntoIterator<Item = u16>) -> ContainerRequest<I>;
144
145    /// Adds a resource ulimit to the container.
146    ///
147    /// # Examples
148    /// ```rust,no_run
149    /// use testcontainers::{GenericImage, ImageExt};
150    ///
151    /// let image = GenericImage::new("image", "tag").with_ulimit("nofile", 65536, Some(65536));
152    /// ```
153    fn with_ulimit(self, name: &str, soft: i64, hard: Option<i64>) -> ContainerRequest<I>;
154
155    /// Sets the container to run in privileged mode.
156    fn with_privileged(self, privileged: bool) -> ContainerRequest<I>;
157
158    /// Adds the capabilities to the container
159    fn with_cap_add(self, capability: impl Into<String>) -> ContainerRequest<I>;
160
161    /// Drops the capabilities from the container's capabilities
162    fn with_cap_drop(self, capability: impl Into<String>) -> ContainerRequest<I>;
163
164    /// cgroup namespace mode for the container. Possible values are:
165    /// - [`CgroupnsMode::Private`]: the container runs in its own private cgroup namespace
166    /// - [`CgroupnsMode::Host`]: use the host system's cgroup namespace
167    ///
168    /// If not specified, the daemon default is used, which can either be `\"private\"` or `\"host\"`, depending on daemon version, kernel support and configuration.
169    fn with_cgroupns_mode(self, cgroupns_mode: CgroupnsMode) -> ContainerRequest<I>;
170
171    /// Sets the usernamespace mode for the container when usernamespace remapping option is enabled.
172    fn with_userns_mode(self, userns_mode: &str) -> ContainerRequest<I>;
173
174    /// Sets the shared memory size in bytes
175    fn with_shm_size(self, bytes: u64) -> ContainerRequest<I>;
176
177    /// Sets the startup timeout for the container. The default is 60 seconds.
178    fn with_startup_timeout(self, timeout: Duration) -> ContainerRequest<I>;
179
180    /// Sets the working directory. The default is defined by the underlying image, which in turn may default to `/`.
181    fn with_working_dir(self, working_dir: impl Into<String>) -> ContainerRequest<I>;
182
183    /// Adds the log consumer to the container.
184    ///
185    /// Allows to follow the container logs for the whole lifecycle of the container, starting from the creation.
186    fn with_log_consumer(self, log_consumer: impl LogConsumer + 'static) -> ContainerRequest<I>;
187
188    /// Flag the container as being exempt from the default `testcontainers` remove-on-drop lifecycle,
189    /// indicating that the container should be kept running, and that executions with the same configuration
190    /// reuse it instead of starting a "fresh" container instance.
191    ///
192    /// **NOTE:** Reusable Containers is an experimental feature, and its behavior is therefore subject
193    /// to change. Containers marked as `reuse` **_will not_** be stopped or cleaned up when their associated
194    /// `Container` or `ContainerAsync` is dropped.
195    #[cfg(feature = "reusable-containers")]
196    fn with_reuse(self, reuse: ReuseDirective) -> ContainerRequest<I>;
197
198    /// Sets the user that commands are run as inside the container.
199    fn with_user(self, user: impl Into<String>) -> ContainerRequest<I>;
200
201    /// Sets the container's root filesystem to be mounted as read-only
202    fn with_readonly_rootfs(self, readonly_rootfs: bool) -> ContainerRequest<I>;
203
204    /// Sets security options for the container
205    fn with_security_opt(self, security_opt: impl Into<String>) -> ContainerRequest<I>;
206
207    /// Overrides ready conditions.
208    ///
209    /// There is no guarantee that the specified ready conditions for an image would result
210    /// in a running container. Users of this API are advised to use this at their own risk.
211    fn with_ready_conditions(self, ready_conditions: Vec<WaitFor>) -> ContainerRequest<I>;
212
213    /// Sets a custom health check for the container.
214    ///
215    /// This will override any `HEALTHCHECK` instruction defined in the image.
216    /// See [`Healthcheck`] for more details on how to build a health check.
217    ///
218    /// # Example
219    ///
220    /// ```rust,no_run
221    /// use testcontainers::{core::{Healthcheck, WaitFor}, GenericImage, ImageExt};
222    /// use std::time::Duration;
223    ///
224    /// let image = GenericImage::new("mysql", "8.0")
225    ///     .with_wait_for(WaitFor::healthcheck())
226    ///     .with_health_check(
227    ///         Healthcheck::cmd(["mysqladmin", "ping", "-h", "localhost", "-u", "root", "-proot"])
228    ///             .with_interval(Duration::from_secs(2))
229    ///             .with_timeout(Duration::from_secs(1))
230    ///             .with_retries(5)
231    ///     );
232    /// ```
233    fn with_health_check(self, health_check: Healthcheck) -> ContainerRequest<I>;
234
235    /// Injects device requests into the container request.
236    ///
237    /// This allows, for instance, exposing the underlying host's GPU:
238    /// https://docs.docker.com/compose/how-tos/gpu-support/#example-of-a-compose-file-for-running-a-service-with-access-to-1-gpu-device
239    ///
240    /// This brings in 2 requirements:
241    ///
242    /// - The host must have [NVIDIA container toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) installed.
243    /// - The image must have [NVIDIA drivers](https://www.nvidia.com/en-us/drivers/) installed.
244    ///
245    /// # Example
246    ///
247    /// ```rust,no_run
248    /// use testcontainers::{GenericImage, ImageExt as _, bollard::models::DeviceRequest};
249    ///
250    /// let device_request = DeviceRequest {
251    ///     driver: Some(String::from("nvidia")),
252    ///     count: Some(-1), // expose all
253    ///     capabilities: Some(vec![vec![String::from("gpu")]]),
254    ///     device_ids: None,
255    ///     options: None,
256    /// };
257    ///
258    /// let image = GenericImage::new("ubuntu", "24.04")
259    ///     .with_device_requests(vec![device_request]);
260    /// ```
261    #[cfg(feature = "device-requests")]
262    fn with_device_requests(self, device_requests: Vec<DeviceRequest>) -> ContainerRequest<I>;
263}
264
265/// Implements the [`ImageExt`] trait for the every type that can be converted into a [`ContainerRequest`].
266impl<RI: Into<ContainerRequest<I>>, I: Image> ImageExt<I> for RI {
267    fn with_cmd(self, cmd: impl IntoIterator<Item = impl Into<String>>) -> ContainerRequest<I> {
268        let container_req = self.into();
269        ContainerRequest {
270            overridden_cmd: cmd.into_iter().map(Into::into).collect(),
271            ..container_req
272        }
273    }
274
275    fn with_name(self, name: impl Into<String>) -> ContainerRequest<I> {
276        let container_req = self.into();
277        ContainerRequest {
278            image_name: Some(name.into()),
279            ..container_req
280        }
281    }
282
283    fn with_tag(self, tag: impl Into<String>) -> ContainerRequest<I> {
284        let container_req = self.into();
285        ContainerRequest {
286            image_tag: Some(tag.into()),
287            ..container_req
288        }
289    }
290
291    fn with_container_name(self, name: impl Into<String>) -> ContainerRequest<I> {
292        let container_req = self.into();
293
294        ContainerRequest {
295            container_name: Some(name.into()),
296            ..container_req
297        }
298    }
299
300    fn with_platform(self, platform: impl Into<String>) -> ContainerRequest<I> {
301        let container_req = self.into();
302
303        ContainerRequest {
304            platform: Some(platform.into()),
305            ..container_req
306        }
307    }
308
309    fn with_network(self, network: impl Into<String>) -> ContainerRequest<I> {
310        let container_req = self.into();
311        ContainerRequest {
312            network: Some(network.into()),
313            ..container_req
314        }
315    }
316
317    fn with_label(self, key: impl Into<String>, value: impl Into<String>) -> ContainerRequest<I> {
318        let mut container_req = self.into();
319
320        container_req.labels.insert(key.into(), value.into());
321
322        container_req
323    }
324
325    fn with_labels(
326        self,
327        labels: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
328    ) -> ContainerRequest<I> {
329        let mut container_req = self.into();
330
331        container_req.labels.extend(
332            labels
333                .into_iter()
334                .map(|(key, value)| (key.into(), value.into())),
335        );
336
337        container_req
338    }
339
340    fn with_env_var(
341        self,
342        name: impl Into<String>,
343        value: impl Into<String>,
344    ) -> ContainerRequest<I> {
345        let mut container_req = self.into();
346        container_req.env_vars.insert(name.into(), value.into());
347        container_req
348    }
349
350    fn with_host(self, key: impl Into<String>, value: impl Into<Host>) -> ContainerRequest<I> {
351        let mut container_req = self.into();
352        container_req.hosts.insert(key.into(), value.into());
353        container_req
354    }
355
356    fn with_hostname(self, hostname: impl Into<String>) -> ContainerRequest<I> {
357        let mut container_req = self.into();
358        container_req.hostname = Some(hostname.into());
359        container_req
360    }
361
362    fn with_mount(self, mount: impl Into<Mount>) -> ContainerRequest<I> {
363        let mut container_req = self.into();
364        container_req.mounts.push(mount.into());
365        container_req
366    }
367
368    fn with_copy_to(
369        self,
370        target: impl Into<String>,
371        source: impl Into<CopyDataSource>,
372    ) -> ContainerRequest<I> {
373        let mut container_req = self.into();
374        let target: String = target.into();
375        container_req
376            .copy_to_sources
377            .push(CopyToContainer::new(source, target));
378        container_req
379    }
380
381    fn with_mapped_port(
382        self,
383        host_port: u16,
384        container_port: ContainerPort,
385    ) -> ContainerRequest<I> {
386        let container_req = self.into();
387        let mut ports = container_req.ports.unwrap_or_default();
388        ports.push(PortMapping::new(host_port, container_port));
389
390        ContainerRequest {
391            ports: Some(ports),
392            ..container_req
393        }
394    }
395
396    #[cfg(feature = "host-port-exposure")]
397    fn with_exposed_host_port(self, port: u16) -> ContainerRequest<I> {
398        self.with_exposed_host_ports([port])
399    }
400
401    #[cfg(feature = "host-port-exposure")]
402    fn with_exposed_host_ports(self, ports: impl IntoIterator<Item = u16>) -> ContainerRequest<I> {
403        let mut container_req = self.into();
404        let exposures = container_req
405            .host_port_exposures
406            .get_or_insert_with(Vec::new);
407
408        exposures.extend(ports);
409        exposures.sort_unstable();
410        exposures.dedup();
411
412        container_req
413    }
414
415    fn with_ulimit(self, name: &str, soft: i64, hard: Option<i64>) -> ContainerRequest<I> {
416        let container_req = self.into();
417        let mut ulimits = container_req.ulimits.unwrap_or_default();
418        ulimits.push(ResourcesUlimits {
419            name: Some(name.into()),
420            soft: Some(soft),
421            hard,
422        });
423
424        ContainerRequest {
425            ulimits: Some(ulimits),
426            ..container_req
427        }
428    }
429
430    fn with_privileged(self, privileged: bool) -> ContainerRequest<I> {
431        let container_req = self.into();
432        ContainerRequest {
433            privileged,
434            ..container_req
435        }
436    }
437
438    fn with_cap_add(self, capability: impl Into<String>) -> ContainerRequest<I> {
439        let mut container_req = self.into();
440        container_req
441            .cap_add
442            .get_or_insert_with(Vec::new)
443            .push(capability.into());
444
445        container_req
446    }
447
448    fn with_cap_drop(self, capability: impl Into<String>) -> ContainerRequest<I> {
449        let mut container_req = self.into();
450        container_req
451            .cap_drop
452            .get_or_insert_with(Vec::new)
453            .push(capability.into());
454
455        container_req
456    }
457
458    fn with_cgroupns_mode(self, cgroupns_mode: CgroupnsMode) -> ContainerRequest<I> {
459        let container_req = self.into();
460        ContainerRequest {
461            cgroupns_mode: Some(cgroupns_mode),
462            ..container_req
463        }
464    }
465
466    fn with_userns_mode(self, userns_mode: &str) -> ContainerRequest<I> {
467        let container_req = self.into();
468        ContainerRequest {
469            userns_mode: Some(String::from(userns_mode)),
470            ..container_req
471        }
472    }
473
474    fn with_shm_size(self, bytes: u64) -> ContainerRequest<I> {
475        let container_req = self.into();
476        ContainerRequest {
477            shm_size: Some(bytes),
478            ..container_req
479        }
480    }
481
482    fn with_startup_timeout(self, timeout: Duration) -> ContainerRequest<I> {
483        let container_req = self.into();
484        ContainerRequest {
485            startup_timeout: Some(timeout),
486            ..container_req
487        }
488    }
489
490    fn with_working_dir(self, working_dir: impl Into<String>) -> ContainerRequest<I> {
491        let container_req = self.into();
492        ContainerRequest {
493            working_dir: Some(working_dir.into()),
494            ..container_req
495        }
496    }
497
498    fn with_log_consumer(self, log_consumer: impl LogConsumer + 'static) -> ContainerRequest<I> {
499        let mut container_req = self.into();
500        container_req.log_consumers.push(Box::new(log_consumer));
501        container_req
502    }
503
504    #[cfg(feature = "reusable-containers")]
505    fn with_reuse(self, reuse: ReuseDirective) -> ContainerRequest<I> {
506        ContainerRequest {
507            reuse,
508            ..self.into()
509        }
510    }
511
512    fn with_user(self, user: impl Into<String>) -> ContainerRequest<I> {
513        let container_req = self.into();
514        ContainerRequest {
515            user: Some(user.into()),
516            ..container_req
517        }
518    }
519
520    fn with_readonly_rootfs(self, readonly_rootfs: bool) -> ContainerRequest<I> {
521        let container_req = self.into();
522        ContainerRequest {
523            readonly_rootfs,
524            ..container_req
525        }
526    }
527
528    fn with_security_opt(self, security_opt: impl Into<String>) -> ContainerRequest<I> {
529        let mut container_req = self.into();
530        container_req
531            .security_opts
532            .get_or_insert_with(Vec::new)
533            .push(security_opt.into());
534
535        container_req
536    }
537
538    fn with_ready_conditions(self, ready_conditions: Vec<WaitFor>) -> ContainerRequest<I> {
539        let mut container_req = self.into();
540        container_req.ready_conditions = Some(ready_conditions);
541        container_req
542    }
543
544    fn with_health_check(self, health_check: Healthcheck) -> ContainerRequest<I> {
545        let mut container_req = self.into();
546        container_req.health_check = Some(health_check);
547        container_req
548    }
549
550    #[cfg(feature = "device-requests")]
551    fn with_device_requests(self, device_requests: Vec<DeviceRequest>) -> ContainerRequest<I> {
552        let container_req = self.into();
553        ContainerRequest {
554            device_requests: Some(device_requests),
555            ..container_req
556        }
557    }
558}
559
560#[cfg(all(test, feature = "host-port-exposure"))]
561mod tests {
562    use super::*;
563    use crate::images::generic::GenericImage;
564
565    #[test]
566    fn test_with_exposed_host_port_single() {
567        let image = GenericImage::new("test", "latest");
568        let request = image.with_exposed_host_port(8080);
569
570        assert_eq!(request.host_port_exposures, Some(vec![8080]));
571    }
572
573    #[test]
574    fn test_with_exposed_host_ports_multiple() {
575        let image = GenericImage::new("test", "latest");
576        let request = image.with_exposed_host_ports([8080, 9090, 3000]);
577
578        assert_eq!(request.host_port_exposures, Some(vec![3000, 8080, 9090]));
579    }
580
581    #[test]
582    fn test_with_exposed_host_ports_deduplication() {
583        let image = GenericImage::new("test", "latest");
584        let request = image.with_exposed_host_ports([8080, 9090, 8080, 3000, 9090]);
585
586        assert_eq!(request.host_port_exposures, Some(vec![3000, 8080, 9090]));
587    }
588
589    #[test]
590    fn test_with_exposed_host_ports_empty() {
591        let image = GenericImage::new("test", "latest");
592        let request = image.with_exposed_host_ports([]);
593
594        assert_eq!(request.host_port_exposures, Some(vec![]));
595    }
596
597    #[test]
598    fn test_with_exposed_host_ports_chaining() {
599        let image = GenericImage::new("test", "latest");
600        let request = image
601            .with_exposed_host_port(8080)
602            .with_exposed_host_ports([9090, 3000]);
603
604        assert_eq!(request.host_port_exposures, Some(vec![3000, 8080, 9090]));
605    }
606
607    #[test]
608    fn test_with_exposed_host_ports_preserves_existing() {
609        let image = GenericImage::new("test", "latest");
610        let request = image.with_exposed_host_port(8080);
611
612        // The first call already set host_port_exposures to Some(vec![8080])
613        // Now we add more ports
614        let request = request.with_exposed_host_ports([9090, 3000]);
615
616        // The result should include all ports: 8080 (from first call), 9090, 3000 (from second call)
617        assert_eq!(request.host_port_exposures, Some(vec![3000, 8080, 9090]));
618    }
619}
testcontainers/core/image/image_ext.rs

testcontainers/core/image/
image_ext.rs
