testcontainers/core/wait/
mod.rs

1use std::{env::var, fmt::Debug, time::Duration};
2
3pub use exit_strategy::ExitWaitStrategy;
4pub use health_strategy::HealthWaitStrategy;
5#[cfg(feature = "http_wait")]
6#[cfg_attr(docsrs, doc(cfg(feature = "http_wait")))]
7pub use http_strategy::HttpWaitStrategy;
8pub use log_strategy::LogWaitStrategy;
9
10use crate::{
11    core::{client::Client, logs::LogSource},
12    ContainerAsync, Image,
13};
14
15pub(crate) mod cmd_wait;
16pub(crate) mod exit_strategy;
17pub(crate) mod health_strategy;
18#[cfg(feature = "http_wait")]
19pub(crate) mod http_strategy;
20pub(crate) mod log_strategy;
21
22pub(crate) trait WaitStrategy {
23    async fn wait_until_ready<I: Image>(
24        self,
25        client: &Client,
26        container: &ContainerAsync<I>,
27    ) -> crate::core::error::Result<()>;
28}
29
30/// Represents a condition that needs to be met before a container is considered ready.
31#[derive(Debug, Clone)]
32pub enum WaitFor {
33    /// An empty condition. Useful for default cases or fallbacks.
34    Nothing,
35    /// Wait for a certain message to appear in the container's logs.
36    Log(LogWaitStrategy),
37    /// Wait for a certain amount of time.
38    Duration { length: Duration },
39    /// Wait for the container's status to become `healthy`.
40    Healthcheck(HealthWaitStrategy),
41    /// Wait for a certain HTTP response.
42    #[cfg(feature = "http_wait")]
43    #[cfg_attr(docsrs, doc(cfg(feature = "http_wait")))]
44    Http(HttpWaitStrategy),
45    /// Wait for the container to exit.
46    Exit(ExitWaitStrategy),
47}
48
49impl WaitFor {
50    /// Wait for the message to appear on the container's stdout.
51    pub fn message_on_stdout(message: impl AsRef<[u8]>) -> WaitFor {
52        Self::log(LogWaitStrategy::new(LogSource::StdOut, message))
53    }
54
55    /// Wait for the message to appear on the container's stderr.
56    pub fn message_on_stderr(message: impl AsRef<[u8]>) -> WaitFor {
57        Self::log(LogWaitStrategy::new(LogSource::StdErr, message))
58    }
59
60    /// Wait for the message to appear on the container's stdout.
61    pub fn log(log_strategy: LogWaitStrategy) -> WaitFor {
62        WaitFor::Log(log_strategy)
63    }
64
65    /// Wait for the container to become healthy.
66    ///
67    /// If you need to customize polling interval, use [`HealthWaitStrategy::with_poll_interval`]
68    /// and create the strategy [`WaitFor::Healthcheck`] manually.
69    pub fn healthcheck() -> WaitFor {
70        WaitFor::Healthcheck(HealthWaitStrategy::default())
71    }
72
73    /// Wait for a certain HTTP response.
74    #[cfg(feature = "http_wait")]
75    #[cfg_attr(docsrs, doc(cfg(feature = "http_wait")))]
76    pub fn http(http_strategy: HttpWaitStrategy) -> WaitFor {
77        WaitFor::Http(http_strategy)
78    }
79
80    /// Wait for the container to exit.
81    pub fn exit(exit_strategy: ExitWaitStrategy) -> WaitFor {
82        WaitFor::Exit(exit_strategy)
83    }
84
85    /// Wait for a certain amount of seconds.
86    ///
87    /// Generally, it's not recommended to use this method, as it's better to wait for a specific condition to be met.
88    pub fn seconds(length: u64) -> WaitFor {
89        WaitFor::Duration {
90            length: Duration::from_secs(length),
91        }
92    }
93
94    /// Wait for a certain amount of millis.
95    ///
96    /// Generally, it's not recommended to use this method, as it's better to wait for a specific condition to be met.
97    pub fn millis(length: u64) -> WaitFor {
98        WaitFor::Duration {
99            length: Duration::from_millis(length),
100        }
101    }
102
103    /// Wait for a certain amount of millis specified in the environment variable.
104    ///
105    /// Generally, it's not recommended to use this method, as it's better to wait for a specific condition to be met.
106    pub fn millis_in_env_var(name: &'static str) -> WaitFor {
107        let additional_sleep_period = var(name).map(|value| value.parse());
108
109        (|| {
110            let length = additional_sleep_period.ok()?.ok()?;
111
112            Some(WaitFor::Duration {
113                length: Duration::from_millis(length),
114            })
115        })()
116        .unwrap_or(WaitFor::Nothing)
117    }
118}
119
120#[cfg(feature = "http_wait")]
121#[cfg_attr(docsrs, doc(cfg(feature = "http_wait")))]
122impl From<HttpWaitStrategy> for WaitFor {
123    fn from(value: HttpWaitStrategy) -> Self {
124        Self::Http(value)
125    }
126}
127
128impl WaitStrategy for WaitFor {
129    async fn wait_until_ready<I: Image>(
130        self,
131        client: &Client,
132        container: &ContainerAsync<I>,
133    ) -> crate::core::error::Result<()> {
134        match self {
135            WaitFor::Log(strategy) => strategy.wait_until_ready(client, container).await?,
136            WaitFor::Duration { length } => {
137                tokio::time::sleep(length).await;
138            }
139            WaitFor::Healthcheck(strategy) => {
140                strategy.wait_until_ready(client, container).await?;
141            }
142            #[cfg(feature = "http_wait")]
143            WaitFor::Http(strategy) => {
144                strategy.wait_until_ready(client, container).await?;
145            }
146            WaitFor::Exit(strategy) => {
147                strategy.wait_until_ready(client, container).await?;
148            }
149            WaitFor::Nothing => {}
150        }
151        Ok(())
152    }
153}