diff --git a/testcontainers/src/core/error.rs b/testcontainers/src/core/error.rs index 79340152..1957b53c 100644 --- a/testcontainers/src/core/error.rs +++ b/testcontainers/src/core/error.rs @@ -62,6 +62,8 @@ pub enum WaitContainerError { Unhealthy, #[error("container startup timeout")] StartupTimeout, + #[error("container exited with unexpected code: expected {expected}, actual {actual:?}")] + UnexpectedExitCode { expected: i64, actual: Option }, } impl TestcontainersError { diff --git a/testcontainers/src/core/wait/exit_strategy.rs b/testcontainers/src/core/wait/exit_strategy.rs new file mode 100644 index 00000000..eed8f3c8 --- /dev/null +++ b/testcontainers/src/core/wait/exit_strategy.rs @@ -0,0 +1,76 @@ +use std::time::Duration; + +use crate::{ + core::{client::Client, error::WaitContainerError, wait::WaitStrategy}, + ContainerAsync, Image, +}; + +#[derive(Debug, Clone)] +pub struct ExitWaitStrategy { + expected_code: Option, + poll_interval: Duration, +} + +impl ExitWaitStrategy { + /// Create a new `ExitWaitStrategy` with default settings. + pub fn new() -> Self { + Self { + expected_code: None, + poll_interval: Duration::from_millis(100), + } + } + + /// Set the poll interval for checking the container's status. + pub fn with_poll_interval(mut self, poll_interval: Duration) -> Self { + self.poll_interval = poll_interval; + self + } + + /// Set the expected exit code of the container. + pub fn with_exit_code(mut self, expected_code: i64) -> Self { + self.expected_code = Some(expected_code); + self + } +} + +impl WaitStrategy for ExitWaitStrategy { + async fn wait_until_ready( + self, + client: &Client, + container: &ContainerAsync, + ) -> crate::core::error::Result<()> { + loop { + let container_state = client + .inspect(container.id()) + .await? + .state + .ok_or(WaitContainerError::StateUnavailable)?; + + let is_running = container_state.running.unwrap_or_default(); + + if is_running { + tokio::time::sleep(self.poll_interval).await; + continue; + } + + if let Some(expected_code) = self.expected_code { + let exit_code = container_state.exit_code; + if exit_code != Some(expected_code) { + return Err(WaitContainerError::UnexpectedExitCode { + expected: expected_code, + actual: exit_code, + } + .into()); + } + } + break; + } + Ok(()) + } +} + +impl Default for ExitWaitStrategy { + fn default() -> Self { + Self::new() + } +} diff --git a/testcontainers/src/core/wait/mod.rs b/testcontainers/src/core/wait/mod.rs index 603b17e9..be233235 100644 --- a/testcontainers/src/core/wait/mod.rs +++ b/testcontainers/src/core/wait/mod.rs @@ -1,5 +1,6 @@ use std::{env::var, fmt::Debug, time::Duration}; +pub use exit_strategy::ExitWaitStrategy; pub use health_strategy::HealthWaitStrategy; pub use http_strategy::HttpWaitStrategy; pub use log_strategy::LogWaitStrategy; @@ -10,6 +11,7 @@ use crate::{ }; pub(crate) mod cmd_wait; +pub(crate) mod exit_strategy; pub(crate) mod health_strategy; pub(crate) mod http_strategy; pub(crate) mod log_strategy; @@ -35,6 +37,8 @@ pub enum WaitFor { Healthcheck(HealthWaitStrategy), /// Wait for a certain HTTP response. Http(HttpWaitStrategy), + /// Wait for the container to exit. + Exit(ExitWaitStrategy), } impl WaitFor { @@ -66,6 +70,11 @@ impl WaitFor { WaitFor::Http(http_strategy) } + /// Wait for the container to exit. + pub fn exit(exit_strategy: ExitWaitStrategy) -> WaitFor { + WaitFor::Exit(exit_strategy) + } + /// Wait for a certain amount of seconds. /// /// Generally, it's not recommended to use this method, as it's better to wait for a specific condition to be met. @@ -124,6 +133,9 @@ impl WaitStrategy for WaitFor { WaitFor::Http(strategy) => { strategy.wait_until_ready(client, container).await?; } + WaitFor::Exit(strategy) => { + strategy.wait_until_ready(client, container).await?; + } WaitFor::Nothing => {} } Ok(()) diff --git a/testcontainers/tests/async_runner.rs b/testcontainers/tests/async_runner.rs index 9ff13b1f..d221e0ad 100644 --- a/testcontainers/tests/async_runner.rs +++ b/testcontainers/tests/async_runner.rs @@ -5,7 +5,7 @@ use reqwest::StatusCode; use testcontainers::{ core::{ logs::{consumer::logging_consumer::LoggingConsumer, LogFrame}, - wait::{HttpWaitStrategy, LogWaitStrategy}, + wait::{ExitWaitStrategy, HttpWaitStrategy, LogWaitStrategy}, CmdWaitFor, ExecCommand, IntoContainerPort, WaitFor, }, runners::AsyncRunner, @@ -26,7 +26,10 @@ impl Image for HelloWorld { } fn ready_conditions(&self) -> Vec { - vec![WaitFor::message_on_stdout("Hello from Docker!")] + vec![ + WaitFor::message_on_stdout("Hello from Docker!"), + WaitFor::exit(ExitWaitStrategy::new().with_exit_code(0)), + ] } }