From edf5f14d59e91932bc030c8a6e23ab75ef2ef2ba Mon Sep 17 00:00:00 2001 From: Joseph Sutton Date: Sun, 30 Mar 2025 14:30:06 -0500 Subject: [PATCH 1/6] feat: added CommandStrategy with SuccessfulCommand enum --- testcontainers/src/core/image/exec.rs | 2 +- .../src/core/wait/command_strategy.rs | 110 ++++++++++++++++++ testcontainers/src/core/wait/mod.rs | 9 ++ testcontainers/tests/async_runner.rs | 29 ++++- 4 files changed, 148 insertions(+), 2 deletions(-) create mode 100644 testcontainers/src/core/wait/command_strategy.rs 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..0844f353 --- /dev/null +++ b/testcontainers/src/core/wait/command_strategy.rs @@ -0,0 +1,110 @@ +use std::time::Duration; + +use crate::{ + core::{client::Client, error::WaitContainerError, wait::WaitStrategy, ExecCommand}, + ContainerAsync, Image, +}; + +#[derive(Debug, Clone)] +pub struct CommandStrategy { + expected_code: i64, + 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(), + expected_code: 0, + 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<()> { + 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 self.fail_fast && exit_code != Some(self.expected_code) { + return Err(WaitContainerError::UnexpectedExitCode { + expected: self.expected_code, + actual: exit_code, + } + .into()); + } + + if exit_code == Some(self.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..c6f98a4d 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::{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 successful code. + SuccessfulCommand(CommandStrategy), } impl WaitFor { @@ -146,6 +152,9 @@ impl WaitStrategy for WaitFor { WaitFor::Exit(strategy) => { strategy.wait_until_ready(client, container).await?; } + WaitFor::SuccessfulCommand(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..1397d579 100644 --- a/testcontainers/tests/async_runner.rs +++ b/testcontainers/tests/async_runner.rs @@ -4,7 +4,7 @@ use bollard::Docker; use testcontainers::{ core::{ logs::{consumer::logging_consumer::LoggingConsumer, LogFrame}, - wait::{ExitWaitStrategy, LogWaitStrategy}, + wait::{CommandStrategy, ExitWaitStrategy, LogWaitStrategy}, CmdWaitFor, ExecCommand, WaitFor, }, runners::AsyncRunner, @@ -100,6 +100,33 @@ 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::SuccessfulCommand( + CommandStrategy::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 res = container + .exec(ExecCommand::new(["pg_isready"]).with_cmd_ready_condition(CmdWaitFor::exit())) + .await?; + + // 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(); From 7d449aa5db6f22aca0bab47ef69e4a9d20513ca3 Mon Sep 17 00:00:00 2001 From: Joseph Sutton Date: Sun, 30 Mar 2025 14:46:36 -0500 Subject: [PATCH 2/6] docs: added wait for variant to wait_strategies document --- docs/features/wait_strategies.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/features/wait_strategies.md b/docs/features/wait_strategies.md index 9c312207..e5a6f0f9 100644 --- a/docs/features/wait_strategies.md +++ b/docs/features/wait_strategies.md @@ -14,6 +14,7 @@ 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. +* `SuccessfulCommand` - wait for a given command to exit successfully (exit code 0). [`Image`](https://docs.rs/testcontainers/latest/testcontainers/core/trait.Image.html) implementation is responsible for returning the appropriate `WaitFor` strategies. From 74755023e1696d76f3bf605d9f8f6115bf840fe4 Mon Sep 17 00:00:00 2001 From: Joseph Sutton Date: Sat, 5 Apr 2025 17:23:16 -0500 Subject: [PATCH 3/6] refactor(wait-for): utilize `WaitFor::command(ExecCommand)` --- testcontainers/src/core/wait/mod.rs | 14 +++++++++++--- testcontainers/tests/async_runner.rs | 21 +++++++++++---------- 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/testcontainers/src/core/wait/mod.rs b/testcontainers/src/core/wait/mod.rs index c6f98a4d..9f9b60fe 100644 --- a/testcontainers/src/core/wait/mod.rs +++ b/testcontainers/src/core/wait/mod.rs @@ -13,7 +13,7 @@ use crate::{ ContainerAsync, Image, }; -use super::{CmdWaitFor, ExecCommand}; +use super::{error::WaitContainerError, CmdWaitFor, ExecCommand}; pub(crate) mod cmd_wait; pub(crate) mod command_strategy; @@ -49,7 +49,7 @@ pub enum WaitFor { /// Wait for the container to exit. Exit(ExitWaitStrategy), /// Wait for a certain command to exit with a successful code. - SuccessfulCommand(CommandStrategy), + Command(CommandStrategy), } impl WaitFor { @@ -68,6 +68,14 @@ impl WaitFor { WaitFor::Log(log_strategy) } + /// Wait for the command to execute successfully. + pub fn command(command: ExecCommand) -> WaitFor { + let cmd_strategy = + CommandStrategy::command(command.with_cmd_ready_condition(CmdWaitFor::exit_code(0))); + + WaitFor::Command(cmd_strategy) + } + /// Wait for the container to become healthy. /// /// If you need to customize polling interval, use [`HealthWaitStrategy::with_poll_interval`] @@ -152,7 +160,7 @@ impl WaitStrategy for WaitFor { WaitFor::Exit(strategy) => { strategy.wait_until_ready(client, container).await?; } - WaitFor::SuccessfulCommand(strategy) => { + WaitFor::Command(strategy) => { strategy.wait_until_ready(client, container).await?; } WaitFor::Nothing => {} diff --git a/testcontainers/tests/async_runner.rs b/testcontainers/tests/async_runner.rs index 1397d579..4b97a072 100644 --- a/testcontainers/tests/async_runner.rs +++ b/testcontainers/tests/async_runner.rs @@ -5,7 +5,7 @@ use testcontainers::{ core::{ logs::{consumer::logging_consumer::LoggingConsumer, LogFrame}, wait::{CommandStrategy, ExitWaitStrategy, LogWaitStrategy}, - CmdWaitFor, ExecCommand, WaitFor, + CmdWaitFor, ContainerState, ExecCommand, WaitFor, }, runners::AsyncRunner, GenericImage, Image, ImageExt, @@ -104,11 +104,8 @@ async fn start_containers_in_parallel() -> anyhow::Result<()> { 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::SuccessfulCommand( - CommandStrategy::command( - ExecCommand::new(["pg_isready"]).with_cmd_ready_condition(CmdWaitFor::exit_code(0)), - ), - )); + let image = GenericImage::new("postgres", "latest") + .with_wait_for(WaitFor::command(ExecCommand::new(["pg_isready"]))); let container = image .with_env_var("POSTGRES_USER", "postgres") .with_env_var("POSTGRES_PASSWORD", "postgres") @@ -116,13 +113,17 @@ async fn async_wait_for_successful_command_strategy() -> anyhow::Result<()> { .start() .await?; - let res = container - .exec(ExecCommand::new(["pg_isready"]).with_cmd_ready_condition(CmdWaitFor::exit())) - .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)); + //assert_eq!(res.exit_code().await?, Some(0)); Ok(()) } From 86fa407084d47600650390b65e37b9b45f3c56fd Mon Sep 17 00:00:00 2001 From: Joseph Sutton Date: Sat, 5 Apr 2025 17:27:23 -0500 Subject: [PATCH 4/6] refactor: allow `WaitFor::command()` to be more general than specific --- testcontainers/src/core/wait/mod.rs | 7 +++---- testcontainers/tests/async_runner.rs | 5 +++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/testcontainers/src/core/wait/mod.rs b/testcontainers/src/core/wait/mod.rs index 9f9b60fe..8613bceb 100644 --- a/testcontainers/src/core/wait/mod.rs +++ b/testcontainers/src/core/wait/mod.rs @@ -48,7 +48,7 @@ pub enum WaitFor { Http(HttpWaitStrategy), /// Wait for the container to exit. Exit(ExitWaitStrategy), - /// Wait for a certain command to exit with a successful code. + /// Wait for a certain command to exit with a specific code. Command(CommandStrategy), } @@ -68,10 +68,9 @@ impl WaitFor { WaitFor::Log(log_strategy) } - /// Wait for the command to execute successfully. + /// Wait for the command to execute. pub fn command(command: ExecCommand) -> WaitFor { - let cmd_strategy = - CommandStrategy::command(command.with_cmd_ready_condition(CmdWaitFor::exit_code(0))); + let cmd_strategy = CommandStrategy::command(command); WaitFor::Command(cmd_strategy) } diff --git a/testcontainers/tests/async_runner.rs b/testcontainers/tests/async_runner.rs index 4b97a072..3fde11aa 100644 --- a/testcontainers/tests/async_runner.rs +++ b/testcontainers/tests/async_runner.rs @@ -104,8 +104,9 @@ async fn start_containers_in_parallel() -> anyhow::Result<()> { 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"]))); + 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") From 8a34464a6ce29206f60ca5ce6356d0c2039eae5d Mon Sep 17 00:00:00 2001 From: Joseph Sutton Date: Sat, 5 Apr 2025 17:48:25 -0500 Subject: [PATCH 5/6] docs: update --- docs/features/wait_strategies.md | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/docs/features/wait_strategies.md b/docs/features/wait_strategies.md index e5a6f0f9..8cfdaa43 100644 --- a/docs/features/wait_strategies.md +++ b/docs/features/wait_strategies.md @@ -14,7 +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. -* `SuccessfulCommand` - wait for a given command to exit successfully (exit code 0). +* `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. From cee5d7ec6029c9d5ee9e2d9ee38018bafd9ca6b1 Mon Sep 17 00:00:00 2001 From: Joseph Sutton Date: Sun, 6 Apr 2025 06:23:45 -0500 Subject: [PATCH 6/6] refactor(command-strategy): use exit code from `ExecCommand` --- .../src/core/wait/command_strategy.rs | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/testcontainers/src/core/wait/command_strategy.rs b/testcontainers/src/core/wait/command_strategy.rs index 0844f353..996c25ca 100644 --- a/testcontainers/src/core/wait/command_strategy.rs +++ b/testcontainers/src/core/wait/command_strategy.rs @@ -1,13 +1,14 @@ use std::time::Duration; use crate::{ - core::{client::Client, error::WaitContainerError, wait::WaitStrategy, ExecCommand}, + core::{ + client::Client, error::WaitContainerError, wait::WaitStrategy, CmdWaitFor, ExecCommand, + }, ContainerAsync, Image, }; #[derive(Debug, Clone)] pub struct CommandStrategy { - expected_code: i64, poll_interval: Duration, command: ExecCommand, fail_fast: bool, @@ -18,7 +19,6 @@ impl CommandStrategy { pub fn new() -> Self { Self { command: ExecCommand::default(), - expected_code: 0, poll_interval: Duration::from_millis(100), fail_fast: false, } @@ -56,6 +56,11 @@ impl WaitStrategy for CommandStrategy { 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()) @@ -82,15 +87,17 @@ impl WaitStrategy for CommandStrategy { let exit_code = inspect_result.exit_code; running = inspect_result.running.unwrap_or(false); - if self.fail_fast && exit_code != Some(self.expected_code) { - return Err(WaitContainerError::UnexpectedExitCode { - expected: self.expected_code, - actual: exit_code, + if let Some(code) = expected_code { + if self.fail_fast && exit_code != expected_code { + return Err(WaitContainerError::UnexpectedExitCode { + expected: code, + actual: exit_code, + } + .into()); } - .into()); } - if exit_code == Some(self.expected_code) { + if exit_code == expected_code { return Ok(()); }