omry-integration 0.22.0

use secrecy::{ExposeSecret, SecretString};
use std::borrow::Cow;
use std::path::Path;
use testcontainers::core::{ContainerPort, ContainerRequest, IntoContainerPort, Mount, WaitFor};
use testcontainers::runners::AsyncRunner;
use testcontainers::{ContainerAsync, Image, ImageExt};
use tracing::Level;
use url::Url;

use crate::containers::{ContainerConfig, ContainerLogConsumer, ImageConfig};
use crate::env_keys;
use crate::helpers::{ContainerMount, HostPorts, TempPathWithString};

/// The path to the data directory used by Omry inside the container.
const DATA_PATH_CONTAINER: &str = "/usr/local/etc/flora";

/// This image name assumes the image is stored in the local registry.
const IMAGE_NAME: &str = "localhost:8777/omry-server";

/// Omry server image tag.
const IMAGE_TAG: &str = "latest";

/// The port the Omry server should bind to inside the container for tarpc.
const CONTAINER_PORT_RPC: u16 = 47473;

/// The web/HTTP port the Omry server should bind to inside the container.  
const CONTAINER_PORT_WEB: u16 = 47474;

/// The message the Omry server should print to stdout when it's ready to receive requests.
const READY_MESSAGE: &str = "flora is ready!";

/// Enum to encode plain or secret environment variables
/// for the flora server image.
#[derive(Debug)]
pub enum EnvVar {
    /// Plain (not secret) environment variable.
    Plain {
        /// Key (variable name).
        key: String,

        /// Value.
        value: String,
    },

    /// Secret environment variable, whose value should not be unintentionally exposed.
    Secret {
        /// Key (variable name).
        key: String,

        /// Value (will be stored in [`SecretString`]).
        value: SecretString,
    },
}

impl EnvVar {
    /// Returns this environment variable as a string (key, value) pair.
    fn as_key_value_pair(&self) -> (&str, &str) {
        match self {
            Self::Plain { key, value } => (key, value),
            Self::Secret { key, value } => (key, value.expose_secret()),
        }
    }
}

/// flora image config.
#[allow(clippy::missing_docs_in_private_items)]
#[derive(Debug)]
pub struct FloraConfig {
    auth_token: EnvVar,
    database_url: EnvVar,
    ready_message: EnvVar,
    rpc_port: EnvVar,
    tracing_level: Option<EnvVar>,
    typesense_api_key: EnvVar,
    typesense_url: EnvVar,
    web_port: EnvVar,
}

impl FloraConfig {
    /// Initializes a new flora image config with the given parameters.
    ///
    /// Automatically generates an auth token.
    #[must_use]
    pub fn new(
        auth_token: SecretString,
        typesense_url: &Url,
        typesense_api_key: SecretString,
    ) -> Self {
        let auth_token = EnvVar::Secret {
            key: env_keys::AUTH_TOKEN.to_string(),
            value: auth_token,
        };

        let typesense_api_key = EnvVar::Secret {
            key: env_keys::TYPESENSE_API_KEY.to_string(),
            value: typesense_api_key,
        };
        let typesense_url = EnvVar::Plain {
            key: env_keys::TYPESENSE_URL.to_string(),
            value: typesense_url.to_string(),
        };

        // Hardcode the ready message until we have a reason to let the caller configure it.
        let ready_message = EnvVar::Plain {
            key: env_keys::READY_MESSAGE_KEY.to_string(),
            value: READY_MESSAGE.to_string(),
        };

        let db_path = Path::new(DATA_PATH_CONTAINER)
            .join("flora.db")
            .display()
            .to_string();

        // ditto with database URL.
        let database_url = EnvVar::Plain {
            key: env_keys::DATABASE_URL.into(),
            value: db_path,
        };

        let rpc_port = EnvVar::Plain {
            key: env_keys::SERVER_PORT_RPC.into(),
            value: CONTAINER_PORT_RPC.to_string(),
        };

        let web_port = EnvVar::Plain {
            key: env_keys::SERVER_PORT_WEB.into(),
            value: CONTAINER_PORT_WEB.to_string(),
        };

        Self {
            auth_token,
            database_url,
            ready_message,
            rpc_port,
            typesense_api_key,
            typesense_url,
            tracing_level: None,
            web_port,
        }
    }

    /// Sets the tracing level for the flora RPC server.
    #[must_use]
    pub fn with_tracing_level(mut self, tracing_level: Level) -> Self {
        let tracing_level = EnvVar::Plain {
            key: env_keys::TRACING_LEVEL.into(),
            value: tracing_level.to_string(),
        };
        self.tracing_level = Some(tracing_level);
        self
    }

    /// Returns the values in this config as a vec of environment variables.
    fn into_vec(self) -> Vec<EnvVar> {
        let mut vars = vec![
            self.auth_token,
            self.database_url,
            self.ready_message,
            self.rpc_port,
            self.typesense_url,
            self.typesense_api_key,
            self.web_port,
        ];

        if let Some(tracing_level) = self.tracing_level {
            vars.push(tracing_level);
        }

        vars
    }
}

/// Omry server image for integration tests.
#[allow(clippy::missing_docs_in_private_items)]
#[derive(Debug)]
pub struct FloraImage {
    name: String,
    tag: String,
    env: Vec<EnvVar>,
    mounts: Vec<Mount>,
    container_ports: Vec<ContainerPort>,
    ready_message: String,
    host_ports: Option<HostPorts>,

    // We need to store this on the struct so the temp directory
    // sticks around until the struct is dropped.
    #[allow(dead_code)]
    data_temp_path: Option<TempPathWithString>,
}

impl FloraImage {
    /// Creates a new flora image instance with the given flora server config and
    /// Docker image config.
    ///
    /// ## Errors
    /// Returns [`anyhow::Error`] if we can't create a temporary path for Omry's data.
    pub fn new(flora_config: FloraConfig, image_config: ImageConfig) -> anyhow::Result<Self> {
        let container_mount = ContainerMount::new(DATA_PATH_CONTAINER)?;
        let container_ports = if image_config.should_expose_ports {
            vec![
                ContainerPort::Udp(CONTAINER_PORT_RPC),
                ContainerPort::Tcp(CONTAINER_PORT_WEB),
            ]
        } else {
            vec![]
        };

        let (_, ready_message) = flora_config.ready_message.as_key_value_pair();
        let ready_message = ready_message.into();

        Ok(Self {
            name: IMAGE_NAME.to_string(),
            tag: IMAGE_TAG.to_string(),
            env: flora_config.into_vec(),
            data_temp_path: container_mount.temp_path,
            mounts: vec![container_mount.mount],
            container_ports,
            ready_message,
            host_ports: None,
        })
    }

    /// Configures this container to use the given ports on the host OS.
    fn with_host_ports(mut self, ports: HostPorts) -> ContainerRequest<Self> {
        self.host_ports = Some(ports);
        self.with_mapped_port(ports.rpc, CONTAINER_PORT_RPC.udp())
            .with_mapped_port(ports.web, CONTAINER_PORT_WEB.tcp())
    }

    /// Configures this container to request any free port from the host OS.
    fn with_auto_host_ports(self) -> ContainerRequest<Self> {
        let (host_ports, _dummy_listeners) = HostPorts::auto();
        self.with_host_ports(host_ports)
    }

    /// Returns the host-side port this container is bound to (if any).
    #[must_use]
    pub fn host_ports(&self) -> Option<HostPorts> {
        self.host_ports
    }
}

impl Image for FloraImage {
    fn name(&self) -> &str {
        &self.name
    }

    fn tag(&self) -> &str {
        &self.tag
    }

    fn ready_conditions(&self) -> Vec<WaitFor> {
        vec![WaitFor::message_on_stdout(&self.ready_message)]
    }

    fn env_vars(
        &self,
    ) -> impl IntoIterator<Item = (impl Into<Cow<'_, str>>, impl Into<Cow<'_, str>>)> {
        self.env.iter().map(EnvVar::as_key_value_pair)
    }

    fn mounts(&self) -> impl IntoIterator<Item = &Mount> {
        &self.mounts
    }

    /// Returns the exposed ports for tarpc and web, in that order.
    fn expose_ports(&self) -> &[ContainerPort] {
        &self.container_ports
    }
}

/// Type alias to reduce the type signature.
pub type FloraContainer = ContainerAsync<FloraImage>;

/// Starts an Omry RPC server container for integration tests.
///
/// ## Errors
/// Returns [`anyhow::Error`] if we can't start the Omry server.
pub async fn start_flora_container<L>(
    flora_config: FloraConfig,
    container_config: ContainerConfig<L>,
) -> anyhow::Result<FloraContainer>
where
    L: ContainerLogConsumer,
{
    let flora_container =
        FloraImage::new(flora_config, container_config.image_config)?.with_auto_host_ports();

    let flora_container = container_config.configure_container(flora_container);
    Ok(flora_container.start().await?)
}