Skip to main content

testcontainers_ext/
lib.rs

1#![crate_name = "testcontainers_ext"]
2
3use bollard::{
4    query_parameters::ListContainersOptions, query_parameters::StopContainerOptions,
5    secret::ContainerSummaryStateEnum,
6};
7use std::future::Future;
8use testcontainers::{ContainerRequest, Image, ImageExt, TestcontainersError};
9
10pub trait ImagePruneExistedLabelExt<I>: Sized + ImageExt<I> + Send
11where
12    I: Image,
13{
14    /// Given a scope, a container label, a prune flag, and a force flag,
15    /// this method will prune the container if the prune flag is true.
16    ///
17    /// Example:
18    ///
19    /// ```
20    /// use tokio::runtime::Runtime;
21    /// use testcontainers::{core::{IntoContainerPort, WaitFor}, runners::AsyncRunner, GenericImage, ImageExt};
22    /// use testcontainers_ext::ImagePruneExistedLabelExt;
23    /// use anyhow::Result;
24    ///
25    /// async fn test () -> Result<()> {
26    ///   let container = GenericImage::new("redis", "7.2.4")
27    ///         .with_exposed_port(6379.tcp())
28    ///         .with_wait_for(WaitFor::message_on_stdout("Ready to accept connections"))
29    ///         .with_prune_existed_label("my-project-scope", "redis", true, true).await?
30    ///         .start()
31    ///         .await?;
32    ///    Ok(())
33    /// }
34    ///
35    /// Runtime::new().unwrap().block_on(test()).unwrap();
36    /// ```
37    ///
38    fn with_prune_existed_label(
39        self,
40        scope: &str,
41        container_label: &str,
42        prune: bool,
43        force: bool,
44    ) -> impl Future<Output = Result<ContainerRequest<I>, TestcontainersError>> + Send {
45        use std::collections::HashMap;
46
47        use bollard::query_parameters::PruneContainersOptions;
48        use testcontainers::core::client::docker_client_instance;
49
50        let testcontainers_project_key = format!("{scope}.testcontainers.scope");
51        let testcontainers_container_key = format!("{scope}.testcontainers.container");
52        let testcontainers_prune_key = format!("{scope}.testcontainers.prune");
53
54        async move {
55            if prune {
56                let client = docker_client_instance().await?;
57
58                let mut filters = HashMap::<String, Vec<String>>::new();
59
60                filters.insert(
61                    String::from("label"),
62                    vec![
63                        format!("{testcontainers_prune_key}=true"),
64                        format!("{}={}", testcontainers_project_key, scope),
65                        format!("{}={}", testcontainers_container_key, container_label),
66                    ],
67                );
68
69                if force {
70                    let result = client
71                        .list_containers(Some(ListContainersOptions {
72                            all: false,
73                            filters: Some(filters.clone()),
74                            ..Default::default()
75                        }))
76                        .await
77                        .map_err(|err| TestcontainersError::Other(Box::new(err)))?;
78
79                    let remove_containers = result
80                        .iter()
81                        .filter(|c| {
82                            c.state
83                                .is_some_and(|s| matches!(s, ContainerSummaryStateEnum::RUNNING))
84                        })
85                        .flat_map(|c| c.id.as_deref())
86                        .collect::<Vec<_>>();
87
88                    futures::future::try_join_all(
89                        remove_containers
90                            .iter()
91                            .map(|c| client.stop_container(c, None::<StopContainerOptions>)),
92                    )
93                    .await
94                    .map_err(|error| TestcontainersError::Other(Box::new(error)))?;
95
96                    #[cfg(feature = "tracing")]
97                    if !remove_containers.is_empty() {
98                        tracing::warn!(name = "stop running containers", result = ?remove_containers);
99                    }
100                }
101
102                let _result = client
103                    .prune_containers(Some(PruneContainersOptions {
104                        filters: Some(filters),
105                    }))
106                    .await
107                    .map_err(|err| TestcontainersError::Other(Box::new(err)))?;
108
109                #[cfg(feature = "tracing")]
110                if _result
111                    .containers_deleted
112                    .as_ref()
113                    .is_some_and(|c| !c.is_empty())
114                {
115                    tracing::warn!(name = "prune existed containers", result = ?_result);
116                }
117            }
118
119            let result = self.with_labels([
120                (testcontainers_prune_key, "true"),
121                (testcontainers_project_key, scope),
122                (testcontainers_container_key, container_label),
123            ]);
124
125            Ok(result)
126        }
127    }
128}
129
130impl<R, I> ImagePruneExistedLabelExt<I> for R
131where
132    R: Sized + ImageExt<I> + Send,
133    I: Image,
134{
135}
136
137pub trait ImageDefaultLogConsumerExt<I>: Sized + ImageExt<I>
138where
139    I: Image,
140{
141    /// Given a container, this method will return a container request with a default log consumer.
142    ///
143    /// Example:
144    ///
145    /// ```
146    /// use tokio::runtime::Runtime;
147    /// use testcontainers::{core::{IntoContainerPort, WaitFor}, runners::AsyncRunner, GenericImage, ImageExt};
148    /// use testcontainers_ext::ImageDefaultLogConsumerExt;
149    /// use anyhow::Result;
150    ///
151    /// async fn test () -> Result<()> {
152    ///     let container = GenericImage::new("redis", "7.2.4")
153    ///         .with_exposed_port(6379.tcp())
154    ///         .with_wait_for(WaitFor::message_on_stdout("Ready to accept connections"))
155    ///         .with_default_log_consumer()
156    ///         .start()
157    ///         .await?;
158    ///   Ok(())
159    /// }
160    ///
161    /// Runtime::new().unwrap().block_on(test()).unwrap();
162    /// ```
163    ///
164    fn with_default_log_consumer(self) -> ContainerRequest<I> {
165        use testcontainers::core::logs::consumer::logging_consumer::LoggingConsumer;
166
167        self.with_log_consumer(
168            LoggingConsumer::new()
169                .with_stdout_level(log::Level::Info)
170                .with_stderr_level(log::Level::Error),
171        )
172    }
173}
174
175impl<R, I> ImageDefaultLogConsumerExt<I> for R
176where
177    R: Sized + ImageExt<I>,
178    I: Image,
179{
180}