Skip to content

Commit

Permalink
feat!: introduce ExitWaitStrategy (#684)
Browse files Browse the repository at this point in the history
Allows you to wait for the container to exit, optionally with a specific
exit code.

Closes #676
  • Loading branch information
DDtKey committed Jul 7, 2024
1 parent 9630a6b commit 73792f9
Show file tree
Hide file tree
Showing 4 changed files with 95 additions and 2 deletions.
2 changes: 2 additions & 0 deletions testcontainers/src/core/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<i64> },
}

impl TestcontainersError {
Expand Down
76 changes: 76 additions & 0 deletions testcontainers/src/core/wait/exit_strategy.rs
Original file line number Diff line number Diff line change
@@ -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<i64>,
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<I: Image>(
self,
client: &Client,
container: &ContainerAsync<I>,
) -> 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()
}
}
12 changes: 12 additions & 0 deletions testcontainers/src/core/wait/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand All @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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(())
Expand Down
7 changes: 5 additions & 2 deletions testcontainers/tests/async_runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -26,7 +26,10 @@ impl Image for HelloWorld {
}

fn ready_conditions(&self) -> Vec<WaitFor> {
vec![WaitFor::message_on_stdout("Hello from Docker!")]
vec![
WaitFor::message_on_stdout("Hello from Docker!"),
WaitFor::exit(ExitWaitStrategy::new().with_exit_code(0)),
]
}
}

Expand Down

0 comments on commit 73792f9

Please sign in to comment.