testcontainers 0.27.3

A library for integration-testing against docker containers from within Rust.
use bytes::Bytes;

use super::RawContainer;
use crate::core::{
    client::Client,
    error::WaitContainerError,
    logs::{LogSource, WaitingStreamWrapper},
    wait::WaitStrategy,
};

#[derive(Debug, Clone)]
pub struct LogWaitStrategy {
    source: LogSource,
    message: Bytes,
    times: usize,
}

impl LogWaitStrategy {
    /// Create a new [`LogWaitStrategy`] that waits for the given message to appear in the standard output logs.
    /// Shortcut for `LogWaitStrategy::new(LogSource::StdOut, message)`.
    pub fn stdout(message: impl AsRef<[u8]>) -> Self {
        Self::new(LogSource::StdOut, message)
    }

    /// Create a new [`LogWaitStrategy`] that waits for the given message to appear in the standard error logs.
    /// Shortcut for `LogWaitStrategy::new(LogSource::StdErr, message)`.
    pub fn stderr(message: impl AsRef<[u8]>) -> Self {
        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 {
        Self {
            source,
            message: Bytes::from(message.as_ref().to_vec()),
            times: 1,
        }
    }

    /// Set the number of times the message should appear in the logs.
    pub fn with_times(mut self, times: usize) -> Self {
        self.times = times;
        self
    }
}

impl WaitStrategy for LogWaitStrategy {
    async fn wait_until_ready(
        self,
        client: &Client,
        container: &RawContainer,
    ) -> crate::core::error::Result<()> {
        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)
            .wait_for_message(self.message, self.times)
            .await
            .map_err(WaitContainerError::from)?;

        Ok(())
    }
}