omry_integration/

helpers.rs

//! Helper functions and data structures for integration tests.
use anyhow::{Context, Result};
use fluent_uri::Uri;
use omry_archiving::{Document, Marshal, Record, RecordParams, ToRecord};
use omry_service::iroh::EndpointAddr;
use std::env;
use std::fs::{self, File};
use std::io::prelude::*;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6, TcpListener, ToSocketAddrs};
use std::path::{Path, PathBuf};
use tempfile::{TempDir, tempdir};
use testcontainers::core::Mount;
use tracing::Level;
use tracing_subscriber::FmtSubscriber;
use url::Url;
use uuid::Uuid;

use crate::db::TempDb;

// Env var key for the Docker host.
const DOCKER_HOST_KEY: &str = "DOCKER_HOST";

/// Environment variable key that corresponds to the Cargo workspace directory.
const WORKSPACE_DIR_KEY: &str = "CARGO_WORKSPACE_DIR";

/// The path to the serialized web page used in tests.
const TEST_ARCHIVE_SUBPATH: &str = "test_data/example.postcard";

/// Maximum number of times we'll try to find a free port to bind to.
const FIND_PORT_MAX_RETRIES: u8 = 5;

/// Gets a free port for the given socket from the OS. Returns the listener and the port.
///
/// The listener is returned to deal with TOCTOU, so that the OS
/// doesn't give the port to another service before we bind to it for real
/// wherever we actually need that port in our tests.
fn free_port_from_socket_addr(
    socket: impl ToSocketAddrs,
    max_attempts: u8,
) -> anyhow::Result<(TcpListener, u16)> {
    let mut attempts_left = max_attempts;
    let socket_addrs = socket.to_socket_addrs()?.collect::<Vec<_>>();

    while attempts_left > 0 {
        if let Ok(listener) = TcpListener::bind(socket_addrs.as_slice()) {
            let port = listener.local_addr()?.port();
            return Ok((listener, port));
        }
        attempts_left -= 1;
    }

    anyhow::bail!("Couldn't bind a listener to a port in {max_attempts} attempts.");
}

/// Gets a free IPv4 port from the OS. Returns the listener and the port.
///
/// The listener is returned to account for TOCTOU, so that the OS
/// doesn't release the port for another service before we bind to it for real
/// in our tests.
///
/// ## Panics
/// If the port can't be found within some number of retries (currently 3).
#[must_use]
pub fn free_port_ipv4() -> (TcpListener, u16) {
    let socket = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0);

    // Okay to panic if we can't find a free port, since we can't run integration tests
    // without a port anyway.
    #[allow(clippy::expect_used)]
    free_port_from_socket_addr(socket, FIND_PORT_MAX_RETRIES).expect("Couldn't find a free port.")
}

/// Gets a free IPv6 port from the OS. See [`free_port_ipv4`] for details about the return value.
///
/// ## Panics
/// If the port can't be found within some number of retries (currently 3).
#[must_use]
pub fn free_port_ipv6() -> (TcpListener, u16) {
    let socket = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 0, 0, 0);

    // Okay to panic if we can't find a free port, since we can't run integration tests
    // without a port anyway.
    #[allow(clippy::expect_used)]
    free_port_from_socket_addr(socket, FIND_PORT_MAX_RETRIES).expect("Couldn't find a free port.")
}

/// The Docker host (where the Docker socket is running).
#[derive(Debug, Clone)]
pub(crate) struct DockerHost {
    uri: Uri<String>,
}

impl DockerHost {
    /// Gets the URI for the `DOCKER_HOST`, if any.
    ///
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - `DOCKER_HOST` is not set
    /// - `DOCKER_HOST` cannot be parsed into a [`Uri`].
    pub(crate) fn from_env() -> Result<Self> {
        let uri = env::var(DOCKER_HOST_KEY)
            .map_err(|e| {
                anyhow::anyhow!("Failed to get the docker host from the environment: {e}")
            })?
            .parse()
            .map_err(|e| anyhow::anyhow!("Can't parse the docker host: {e}"))?;

        Ok(Self { uri })
    }

    /// Returns [`true`] if the `DOCKER_HOST` appears to be remote.
    ///
    /// The heuristic is (all must be true):
    ///
    /// - `DOCKER_HOST` env var is set,
    /// - `DOCKER_HOST` URI scheme is `tcp`,
    /// - `DOCKER_HOST` host is pretent
    pub(crate) fn is_remote(&self) -> bool {
        self.uri.scheme().as_str() == "tcp" && self.host().is_some()
    }

    /// Gets the host component of the Docker host (if any).
    ///
    /// Will only be present if the Docker host is remote.
    pub(crate) fn host(&self) -> Option<&str> {
        Some(self.uri.authority()?.host())
    }
}

/// Takes the given subpath, and returns the canonical path within
/// the current Cargo project.
///
/// ## Errors
/// Returns [`anyhow::Error`] if we fail to:
///
/// - Get the path to the current project via `CARGO_WORKSPACE_DIR` env variable
/// - Get the full canonical path from it
pub fn path_in_project(subpath: &str) -> Result<PathBuf> {
    let (_, workspace_dir_path) = env::vars()
        .find(|(k, _)| k == WORKSPACE_DIR_KEY)
        .ok_or_else(|| anyhow::anyhow!("{WORKSPACE_DIR_KEY} is not set"))?;

    let workspace_dir_path = fs::canonicalize(workspace_dir_path)?;
    Ok(workspace_dir_path.join(subpath))
}

/// Loads a document to be used in tests.
///
/// ## Errors
/// Returns [`anyhow::Error`] if we fail to:
///
/// - Get the path to the current project via `CARGO_WORKSPACE_DIR` env variable
/// - Get the full canonical path from it
/// - Read the file from that path
/// - Deserialize the data from that file
pub fn make_test_archive_holder() -> anyhow::Result<Document> {
    let test_archive_path = path_in_project(TEST_ARCHIVE_SUBPATH)?;
    let mut file = File::open(test_archive_path)?;
    let mut buf = vec![];
    file.read_to_end(&mut buf)?;
    let archive_holder = Document::from_bytes(&buf)?;
    Ok(archive_holder)
}

/// Makes an instance of [`RecordParams`] for tests.
///
/// ## Errors
/// Returns [`anyhow::Error`] if the test document can't be loaded.
pub fn make_test_record_params() -> anyhow::Result<RecordParams> {
    let document = make_test_archive_holder()?;
    let url = "https://example.org/".to_string();
    let title = "Example".to_string();
    let client_datetime = "2024-10-11T13:49:46-05:00".to_string();
    let language = Some("en".to_string());

    Ok(RecordParams {
        id: None,
        url,
        title,
        language,
        client_datetime,
        document,
        timestamp_flora: None,
        updated_at: None,
    })
}

/// Makes a [`Record`] instance for tests.
///
/// ## Errors
/// Returns [`anyhow::Error`] if the test record can't be loaded.
pub fn make_test_record() -> anyhow::Result<Record> {
    let incoming_record = Record::new(make_test_record_params()?);
    Ok(incoming_record)
}

/// Gets records that have been inserted into the in-memory SQLite,
/// so that they have unique ids.
///
/// ## Errors
/// Returns [`anyhow::Error`] if any of the following operations fail:
///
/// - We can't create a new [`TempDb`].
/// - We can't insert the test record into SQLite.
/// - We can't get that record back from SQLite.
pub async fn make_test_db_records(how_many: usize) -> anyhow::Result<Vec<impl ToRecord>> {
    let db = TempDb::new().await?;
    let mut db_records = Vec::with_capacity(how_many);
    for _ in 0..how_many {
        let record = make_test_record()?;
        let row_id = db.insert(record.try_into()?).await?;
        let db_record = db.get_record(row_id).await?;
        db_records.push(db_record);
    }
    Ok(db_records)
}

/// Temporary path and its string representation.
#[derive(Debug)]
pub struct TempPathWithString {
    /// The path.
    path: PathBuf,

    /// The string representation of the same path, to make it easier to use
    /// in some functions.
    path_string: String,
}

impl TempPathWithString {
    /// Creates a new [`tempfile::TempDir`], and stores its associated path
    /// and string representation.
    ///
    /// ## Errors
    /// If we can't create a temporary directory or construct the string representation
    /// of its path.
    pub fn new() -> anyhow::Result<Self> {
        let temp_path = tempdir()?.keep();

        // If we can't get a string representation of the temp path, then its display
        // repesentation is unlikely to be useful.
        // So clippy's suggestion isn't helpful here.
        #[allow(clippy::unnecessary_debug_formatting)]
        let temp_path_string = temp_path
            .to_str()
            .ok_or_else(|| anyhow::anyhow!("Couldn't get a path string from {temp_path:?}"))?
            .to_string();

        Ok(Self {
            path: temp_path,
            path_string: temp_path_string,
        })
    }

    /// Returns a reference to this temporary directory's path.
    #[must_use]
    pub fn as_path(&self) -> &Path {
        &self.path
    }

    /// Returns a reference to this temporary path's string representation.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.path_string
    }
}

/// Initializes a tracing subscriber for tests.
///
/// ## Errors
/// Returns [`anyhow::Error`] if
/// we can't set the tracing subscriber as the global default.
pub fn init_tracing_subscriber() -> anyhow::Result<()> {
    let subscriber = FmtSubscriber::builder()
        .with_max_level(Level::INFO)
        .finish();
    tracing::subscriber::set_global_default(subscriber)?;

    Ok(())
}

/// A temporary location for storing traces.
///
/// Used for tests that need to verify certain events happened
/// based on the application traces.
#[derive(Debug)]
pub struct TempTracingLocation {
    /// The temporary directory to store the trace.
    dir: TempDir,

    /// The string representation of the temp tracing path, pre-computed
    /// in [`TempTracingLocation::new`].
    /// Storing the pre-computed string makes it easier for users of this
    /// structure to get a string path without having to do their own error handling.
    dir_path: String,

    /// The UUID7 prefix of the trace files in the temporary directory.
    file_prefix: String,
}

impl TempTracingLocation {
    /// Creates a new temporary directory for trace files, and generates their
    /// UUID7 filename prefix.
    ///
    /// ## Errors
    /// Returns [`anyhow::Error`] if we fail to:
    /// - Create a new temporary directory
    /// - Get the string representation of that directory's path
    pub fn new() -> Result<Self> {
        let dir = TempDir::new()?;
        let dir_path = dir
            .path()
            .to_str()
            .ok_or_else(|| anyhow::anyhow!("Cannot get the string path to the tracing dir."))?
            .to_string();

        let file_prefix = Uuid::now_v7().to_string();

        Ok(Self {
            dir,
            dir_path,
            file_prefix,
        })
    }

    /// Returns a reference to the temporary directory which will hold the trace files.
    #[must_use]
    pub fn dir(&self) -> &TempDir {
        &self.dir
    }

    /// Returns a string representation of the path, corresponding to [`Self::dir`].
    #[must_use]
    pub fn path(&self) -> &str {
        &self.dir_path
    }

    /// The trace filename prefix.
    #[must_use]
    pub fn file_prefix(&self) -> &str {
        &self.file_prefix
    }

    /// Iterates over the temporary directory and finds the first trace file that matches the [`Self::file_prefix`].
    ///
    /// ## Returns
    /// The contents of the found trace file, if any.
    ///
    /// ## Errors
    /// Returns [`anyhow::Error`] if we fail to:
    /// - Read the temporary directory
    /// - Get entries in that directory
    /// - Read the entry matching the given prefix into a string
    pub fn read_first_trace(&self) -> Result<Option<String>> {
        let dir = fs::read_dir(self.dir())?;

        for entry in dir {
            let entry = entry?;
            let file_name = entry.file_name();
            let file_name = file_name.to_string_lossy();

            if file_name.starts_with(self.file_prefix()) {
                let output = fs::read_to_string(entry.path())?;
                return Ok(Some(output));
            }
        }

        Ok(None)
    }
}

/// Host-side ports for the Omry server container to use.
#[derive(Debug, Copy, Clone)]
pub struct HostPorts {
    /// Port for tarpc to bind to.
    pub rpc: u16,

    /// Port for the web server to bind to.
    pub web: u16,
}

impl HostPorts {
    /// Requests any free ports from the host OS.
    ///
    /// Returns dummy listeners for both respective ports to avoid
    /// the ports becoming unavailable due to TOCTOU before
    /// the caller can bind to them.
    #[must_use]
    pub(crate) fn auto() -> (Self, (TcpListener, TcpListener)) {
        let (rpc_listener, rpc) = free_port_ipv6();
        let (web_listener, web) = free_port_ipv6();
        (Self { rpc, web }, (rpc_listener, web_listener))
    }
}

/// A mount for a container.
#[derive(Debug)]
pub struct ContainerMount {
    /// The mount for the container.
    pub mount: Mount,

    /// Temporary path on the host. Used if the Docker host is local.
    pub temp_path: Option<TempPathWithString>,
}

impl ContainerMount {
    /// Creates a mount for the container with the given path.
    ///
    /// If we're using a remote Docker host (set by `USING_REMOTE_DOCKER_HOST_KEY`),
    /// will use [`Mount::tmpfs_mount`].
    /// This assumes the remote Docker host is running on Linux (and will fail otherwise).
    ///
    /// Otherwise, will use [`Mount::bind_mount`].
    ///
    /// # Errors
    ///
    /// Returns [`anyhow::Error`] if the Docker host is local,
    /// and we can't create a temporary path on the host.
    pub fn new(container_path: &str) -> Result<Self> {
        if let Ok(docker_host) = DockerHost::from_env()
            && docker_host.is_remote()
        {
            let mount = Mount::tmpfs_mount(container_path);
            return Ok(Self {
                mount,
                temp_path: None,
            });
        }

        let temp_path = TempPathWithString::new()?;
        let mount = Mount::bind_mount(temp_path.as_str(), container_path);

        Ok(Self {
            mount,
            temp_path: Some(temp_path),
        })
    }
}

/// Gets the iroh endpoint address for the given server,
/// via its web API.
#[allow(clippy::missing_errors_doc)]
pub async fn endpoint_addr(server_root: &str) -> Result<EndpointAddr> {
    const ENDPOINT_SUBPATH: &str = "/endpoint-addr";
    let url = Url::parse(server_root).context("error parsing the URL")?;
    let url = url.join(ENDPOINT_SUBPATH)?;
    let bytes = reqwest::get(url)
        .await?
        .error_for_status()
        .context("Couldn't get the endpoint_address")?
        .bytes()
        .await
        .context("Couldn't read the bytes")?;
    let endpoint_addr =
        postcard::from_bytes(&bytes).context("Error deserializing the endpoint address")?;
    Ok(endpoint_addr)
}

/// Module for generating an authentication token.
///
/// Inspired by
/// <https://github.com/EffortlessMetrics/uselesskey/blob/52375d7e2987de345006ce50e7831383ddbc0ab4/crates/uselesskey-core-token-shape/src/lib.rs#L70>.
pub mod token {
    use rand::RngExt;
    use rand::distr::Alphanumeric;
    use rand::rngs::ChaCha20Rng;

    /// Generates a new random token with the given length.
    pub fn generate(len: usize) -> String {
        let rng = rand::make_rng::<ChaCha20Rng>();
        rng.sample_iter(Alphanumeric)
            .take(len)
            .map(char::from)
            .collect()
    }
}