diff --git a/testcontainers/src/core/client.rs b/testcontainers/src/core/client.rs index 6878b5b6..92d36c26 100644 --- a/testcontainers/src/core/client.rs +++ b/testcontainers/src/core/client.rs @@ -127,6 +127,11 @@ impl Client { .into_stderr() } + pub(crate) fn both_std_logs(&self, id: &str, follow: bool) -> RawLogStream { + self.logs_stream(id, Some(LogSource::BothStd), follow) + .into_both_std() + } + pub(crate) fn logs(&self, id: &str, follow: bool) -> LogStream { self.logs_stream(id, None, follow) } @@ -266,8 +271,12 @@ impl Client { ) -> LogStream { let options = LogsOptions { follow, - stdout: source_filter.map(LogSource::is_stdout).unwrap_or(true), - stderr: source_filter.map(LogSource::is_stderr).unwrap_or(true), + stdout: source_filter + .map(LogSource::includes_stdout) + .unwrap_or(true), + stderr: source_filter + .map(LogSource::includes_stderr) + .unwrap_or(true), tail: "all".to_owned(), ..Default::default() }; diff --git a/testcontainers/src/core/logs.rs b/testcontainers/src/core/logs.rs index 8eb85bcc..29d481a8 100644 --- a/testcontainers/src/core/logs.rs +++ b/testcontainers/src/core/logs.rs @@ -29,15 +29,16 @@ pub enum WaitLogError { pub enum LogSource { StdOut, StdErr, + BothStd, } impl LogSource { - pub(super) fn is_stdout(self) -> bool { - matches!(self, Self::StdOut) + pub(super) fn includes_stdout(self) -> bool { + matches!(self, Self::StdOut | Self::BothStd) } - pub(super) fn is_stderr(self) -> bool { - matches!(self, Self::StdErr) + pub(super) fn includes_stderr(self) -> bool { + matches!(self, Self::StdErr | Self::BothStd) } } diff --git a/testcontainers/src/core/logs/stream.rs b/testcontainers/src/core/logs/stream.rs index 4299ca67..3343ed55 100644 --- a/testcontainers/src/core/logs/stream.rs +++ b/testcontainers/src/core/logs/stream.rs @@ -6,7 +6,7 @@ use std::{ }; use bytes::Bytes; -use futures::{stream::BoxStream, Stream, StreamExt}; +use futures::{stream::BoxStream, Stream, StreamExt, TryStreamExt}; use tokio_stream::wrappers::UnboundedReceiverStream; use crate::core::logs::LogFrame; @@ -51,6 +51,16 @@ impl LogStream { .boxed() } + /// Log stream with messages from bith stdout and stderr. + pub(crate) fn into_both_std(self) -> RawLogStream { + self.inner + .map_ok(|frame| match frame { + LogFrame::StdErr(bytes) => bytes, + LogFrame::StdOut(bytes) => bytes, + }) + .boxed() + } + /// Splits the log stream into two streams, one for stdout and one for stderr. pub(crate) async fn split(self) -> (RawLogStream, RawLogStream) { let (stdout_tx, stdout_rx) = tokio::sync::mpsc::unbounded_channel(); diff --git a/testcontainers/src/core/wait/log_strategy.rs b/testcontainers/src/core/wait/log_strategy.rs index 6347bdbf..259811c7 100644 --- a/testcontainers/src/core/wait/log_strategy.rs +++ b/testcontainers/src/core/wait/log_strategy.rs @@ -30,6 +30,13 @@ impl LogWaitStrategy { Self::new(LogSource::StdErr, message) } + /// Create a new [`LogWaitStrategy`] that waits for the given message to appear in either + /// standard output logs or standard error logs. + /// Shortcut for `LogWaitStrategy::new(LogSource::BothStd, message)`. + pub fn stdout_or_stderr(message: impl AsRef<[u8]>) -> Self { + Self::new(LogSource::BothStd, message) + } + /// Create a new `LogWaitStrategy` with the given log source and message. /// The message is expected to appear in the logs exactly once by default. pub fn new(source: LogSource, message: impl AsRef<[u8]>) -> Self { @@ -56,6 +63,7 @@ impl WaitStrategy for LogWaitStrategy { let log_stream = match self.source { LogSource::StdOut => client.stdout_logs(container.id(), true), LogSource::StdErr => client.stderr_logs(container.id(), true), + LogSource::BothStd => client.both_std_logs(container.id(), true), }; WaitingStreamWrapper::new(log_stream) diff --git a/testcontainers/src/core/wait/mod.rs b/testcontainers/src/core/wait/mod.rs index d6752a47..3f506170 100644 --- a/testcontainers/src/core/wait/mod.rs +++ b/testcontainers/src/core/wait/mod.rs @@ -57,6 +57,11 @@ impl WaitFor { Self::log(LogWaitStrategy::new(LogSource::StdErr, message)) } + /// Wait for the message to appear on either container's stdout or stderr. + pub fn message_on_either_std(message: impl AsRef<[u8]>) -> WaitFor { + Self::log(LogWaitStrategy::new(LogSource::BothStd, message)) + } + /// Wait for the message to appear on the container's stdout. pub fn log(log_strategy: LogWaitStrategy) -> WaitFor { WaitFor::Log(log_strategy)