diff --git a/docs/features/wait_strategies.md b/docs/features/wait_strategies.md index 9c312207..8cfdaa43 100644 --- a/docs/features/wait_strategies.md +++ b/docs/features/wait_strategies.md @@ -14,6 +14,17 @@ enum with the following variants: * `Healthcheck` - wait for the container to be healthy * `Http` - wait for an HTTP(S) response with predefined conditions (see [`HttpWaitStrategy`](https://docs.rs/testcontainers/latest/testcontainers/core/wait/struct.HttpWaitStrategy.html) for more details) * `Duration` - wait for a specific duration. Usually less preferable and better to combine with other strategies. +* `Command` - wait for a given command to exit successfully (exit code 0). + +## Waiting for a command + +You can wait for a specific command on an image with a specific error code if you wish. For example, let's wait for a Testcontainer with a Postgres image to be ready by checking for the successful exit code `0` of `pg_isready`: + +```rust +let container = GenericImage::new("postgres", "latest").with_wait_for(WaitFor::command( + ExecCommand::new(["pg_isready"]).with_cmd_ready_condition(CmdWaitFor::exit_code(0)), +)); +``` [`Image`](https://docs.rs/testcontainers/latest/testcontainers/core/trait.Image.html) implementation is responsible for returning the appropriate `WaitFor` strategies. diff --git a/testcontainers/src/core/image/exec.rs b/testcontainers/src/core/image/exec.rs index 42208ac7..46b53dbd 100644 --- a/testcontainers/src/core/image/exec.rs +++ b/testcontainers/src/core/image/exec.rs @@ -1,6 +1,6 @@ use crate::core::{CmdWaitFor, WaitFor}; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct ExecCommand { pub(crate) cmd: Vec, pub(crate) cmd_ready_condition: CmdWaitFor, diff --git a/testcontainers/src/core/wait/command_strategy.rs b/testcontainers/src/core/wait/command_strategy.rs new file mode 100644 index 00000000..996c25ca --- /dev/null +++ b/testcontainers/src/core/wait/command_strategy.rs @@ -0,0 +1,117 @@ +use std::time::Duration; + +use crate::{ + core::{ + client::Client, error::WaitContainerError, wait::WaitStrategy, CmdWaitFor, ExecCommand, + }, + ContainerAsync, Image, +}; + +#[derive(Debug, Clone)] +pub struct CommandStrategy { + poll_interval: Duration, + command: ExecCommand, + fail_fast: bool, +} + +impl CommandStrategy { + /// Create a new `CommandStrategy` with default settings. + pub fn new() -> Self { + Self { + command: ExecCommand::default(), + poll_interval: Duration::from_millis(100), + fail_fast: false, + } + } + + /// Creates a new `CommandStrategy` with default settings and a preset command to execute. + pub fn command(command: ExecCommand) -> Self { + CommandStrategy::default().with_exec_command(command) + } + + /// Set the fail fast flag for the strategy, meaning that if the command's first run does not + /// have the expected exit code, the strategy will exit with failure. If the flag is not set, + /// the strategy will continue to poll the container until the expected exit code is reached. + pub fn with_fail_fast(mut self, fail_fast: bool) -> Self { + self.fail_fast = fail_fast; + self + } + + /// Set the command for executing the command on the container. + pub fn with_exec_command(mut self, command: ExecCommand) -> Self { + self.command = command; + self + } + + /// 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 + } +} + +impl WaitStrategy for CommandStrategy { + async fn wait_until_ready( + self, + client: &Client, + container: &ContainerAsync, + ) -> crate::core::error::Result<()> { + let expected_code = match self.command.clone().cmd_ready_condition { + CmdWaitFor::Exit { code } => code, + _ => Some(0), + }; + + 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 { + let exec_result = client + .exec(container.id(), self.command.clone().cmd) + .await?; + + let inspect_result = client.inspect_exec(&exec_result.id).await?; + let mut running = inspect_result.running.unwrap_or(false); + + loop { + if !running { + break; + } + + let inspect_result = client.inspect_exec(&exec_result.id).await?; + let exit_code = inspect_result.exit_code; + running = inspect_result.running.unwrap_or(false); + + if let Some(code) = expected_code { + if self.fail_fast && exit_code != expected_code { + return Err(WaitContainerError::UnexpectedExitCode { + expected: code, + actual: exit_code, + } + .into()); + } + } + + if exit_code == expected_code { + return Ok(()); + } + + tokio::time::sleep(self.poll_interval).await; + } + + continue; + } + } + } +} + +impl Default for CommandStrategy { + fn default() -> Self { + Self::new() + } +} diff --git a/testcontainers/src/core/wait/mod.rs b/testcontainers/src/core/wait/mod.rs index afb258d7..8613bceb 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 command_strategy::CommandStrategy; pub use exit_strategy::ExitWaitStrategy; pub use health_strategy::HealthWaitStrategy; #[cfg(feature = "http_wait")] @@ -12,7 +13,10 @@ use crate::{ ContainerAsync, Image, }; +use super::{error::WaitContainerError, CmdWaitFor, ExecCommand}; + pub(crate) mod cmd_wait; +pub(crate) mod command_strategy; pub(crate) mod exit_strategy; pub(crate) mod health_strategy; #[cfg(feature = "http_wait")] @@ -44,6 +48,8 @@ pub enum WaitFor { Http(HttpWaitStrategy), /// Wait for the container to exit. Exit(ExitWaitStrategy), + /// Wait for a certain command to exit with a specific code. + Command(CommandStrategy), } impl WaitFor { @@ -62,6 +68,13 @@ impl WaitFor { WaitFor::Log(log_strategy) } + /// Wait for the command to execute. + pub fn command(command: ExecCommand) -> WaitFor { + let cmd_strategy = CommandStrategy::command(command); + + WaitFor::Command(cmd_strategy) + } + /// Wait for the container to become healthy. /// /// If you need to customize polling interval, use [`HealthWaitStrategy::with_poll_interval`] @@ -146,6 +159,9 @@ impl WaitStrategy for WaitFor { WaitFor::Exit(strategy) => { strategy.wait_until_ready(client, container).await?; } + WaitFor::Command(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 c0513ce8..3fde11aa 100644 --- a/testcontainers/tests/async_runner.rs +++ b/testcontainers/tests/async_runner.rs @@ -4,8 +4,8 @@ use bollard::Docker; use testcontainers::{ core::{ logs::{consumer::logging_consumer::LoggingConsumer, LogFrame}, - wait::{ExitWaitStrategy, LogWaitStrategy}, - CmdWaitFor, ExecCommand, WaitFor, + wait::{CommandStrategy, ExitWaitStrategy, LogWaitStrategy}, + CmdWaitFor, ContainerState, ExecCommand, WaitFor, }, runners::AsyncRunner, GenericImage, Image, ImageExt, @@ -100,6 +100,35 @@ async fn start_containers_in_parallel() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn async_wait_for_successful_command_strategy() -> anyhow::Result<()> { + let _ = pretty_env_logger::try_init(); + + let image = GenericImage::new("postgres", "latest").with_wait_for(WaitFor::command( + ExecCommand::new(["pg_isready"]).with_cmd_ready_condition(CmdWaitFor::exit_code(0)), + )); + let container = image + .with_env_var("POSTGRES_USER", "postgres") + .with_env_var("POSTGRES_PASSWORD", "postgres") + .with_env_var("POSTGRES_DB", "db") + .start() + .await?; + + let mut out = String::new(); + container.stdout(false).read_to_string(&mut out).await?; + + assert!( + out.contains("server started"), + "stdout must contain 'server started'" + ); + + // if the container.exec exits with 0, then it means the wait_for successful command strategy + // worked + //assert_eq!(res.exit_code().await?, Some(0)); + + Ok(()) +} + #[tokio::test] async fn async_run_exec() -> anyhow::Result<()> { let _ = pretty_env_logger::try_init();