tsunagi/crates/tsunagi-core/src/testing.rs

//! A harness for testing against a real agent.
//!
//! Behind the `testing` feature, because it is only of use to a test and
//! pulls in `tempfile` and a tracing subscriber. It lives here rather than
//! being copied into each crate's `tests/` directory: a protocol plugin in
//! its own crate needs exactly the same harness, and two copies of it would
//! drift.
//!
//! Everything here builds agents that touch nothing outside a temporary
//! directory: loopback only, no relay, no address publication.

// A harness fails a test loudly and at once; that is the whole job. The
// crate forbids these elsewhere for the opposite reason — nothing off the
// network may panic — and this module never sees anything off the network.
#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]

use std::future::Future;
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};

use crate::agent::Event;
use crate::config::{AgentConfig, StoragePaths, TransportPolicy};
use crate::discovery::SharedMemoryDiscovery;
use crate::identity::{NetworkName, NetworkSecret};
use crate::{Agent, Result};
use tempfile::TempDir;
use tokio::sync::broadcast::error::RecvError;

/// Installs a tracing subscriber when `TSUNAGI_TEST_LOG` is set.
///
/// The library never installs a global subscriber itself; tests opt in.
pub fn init_tracing() {
    use std::sync::Once;
    static ONCE: Once = Once::new();
    ONCE.call_once(|| {
        if let Ok(filter) = std::env::var("TSUNAGI_TEST_LOG") {
            let _ = tracing_subscriber::fmt()
                .with_env_filter(tracing_subscriber::EnvFilter::new(filter))
                .with_writer(std::io::stderr)
                .try_init();
        }
    });
}

/// Overall deadline for anything a test waits on.
pub const DEADLINE: Duration = Duration::from_secs(30);

/// How often condition polling re-checks. Never used as the primary sync.
const POLL_INTERVAL: Duration = Duration::from_millis(25);

/// Builds a fully local configuration rooted at `dir`.
///
/// Loopback binding with a dynamic port, no relay, no address lookup and no
/// port mapping, so the suite needs neither the internet nor privileges.
/// Timeouts are shortened so that failure paths finish quickly.
pub fn local_config(dir: &Path) -> AgentConfig {
    let limits = crate::Limits {
        dial_timeout: Duration::from_millis(1500),
        handshake_timeout: Duration::from_secs(5),
        ..Default::default()
    };
    AgentConfig::new(StoragePaths::under(dir))
        .with_transport(TransportPolicy::LocalOnly)
        .with_loopback_bind()
        .with_discovery_interval(Duration::from_millis(150))
        .with_limits(limits)
}

/// Builds a local configuration wired to a shared in-memory discovery table.
pub fn config_with(dir: &Path, discovery: &SharedMemoryDiscovery) -> AgentConfig {
    local_config(dir).with_discovery(Arc::new(discovery.clone()))
}

/// A temporary directory plus the agent running on it.
pub struct TestAgent {
    /// The directory it keeps its state in, removed when this is dropped.
    pub dir: TempDir,
    /// The agent itself.
    pub agent: Agent,
}

impl TestAgent {
    /// Starts an agent on a fresh temporary directory.
    pub async fn spawn(discovery: &SharedMemoryDiscovery) -> Result<Self> {
        init_tracing();
        let dir = TempDir::new().expect("temp dir");
        let agent = Agent::spawn(config_with(dir.path(), discovery)).await?;
        Ok(Self { dir, agent })
    }

    /// Starts an agent with a caller-supplied configuration on a fresh dir.
    pub async fn spawn_with(
        build: impl FnOnce(AgentConfig) -> AgentConfig,
        discovery: &SharedMemoryDiscovery,
    ) -> Result<Self> {
        init_tracing();
        let dir = TempDir::new().expect("temp dir");
        let agent = Agent::spawn(build(config_with(dir.path(), discovery))).await?;
        Ok(Self { dir, agent })
    }

    /// Stops the agent and returns the directory so it can be reopened.
    pub async fn stop(self) -> TempDir {
        self.agent.shutdown().await;
        self.dir
    }
}

/// A network name and a fresh high-entropy secret.
pub fn network(name: &str) -> (NetworkName, NetworkSecret) {
    (
        NetworkName::new(name).expect("valid network name"),
        NetworkSecret::generate(),
    )
}

/// Waits for an event matching `predicate`, under the global deadline.
pub async fn wait_event<T>(
    rx: &mut tokio::sync::broadcast::Receiver<Event>,
    predicate: impl Fn(&Event) -> Option<T>,
) -> T {
    let deadline = Instant::now() + DEADLINE;
    let mut seen: Vec<String> = Vec::new();
    loop {
        let remaining = deadline
            .checked_duration_since(Instant::now())
            .unwrap_or_default();
        assert!(!remaining.is_zero(), "timed out waiting for an event");

        match tokio::time::timeout(remaining, rx.recv()).await {
            Ok(Ok(event)) => {
                if let Some(value) = predicate(&event) {
                    return value;
                }
                if seen.len() < 12 {
                    let label = format!("{event:?}");
                    seen.push(label.chars().take(70).collect());
                }
            }
            Ok(Err(RecvError::Lagged(skipped))) => {
                panic!("event subscriber lagged, missed {skipped} events; first seen: {seen:#?}");
            }
            Ok(Err(RecvError::Closed)) => panic!("event channel closed while waiting"),
            Err(_) => panic!("timed out waiting for an event"),
        }
    }
}

/// Polls an async condition until it returns `Some`, under the global deadline.
pub async fn wait_until<T, F, Fut>(what: &str, mut probe: F) -> T
where
    F: FnMut() -> Fut,
    Fut: Future<Output = Option<T>>,
{
    let deadline = Instant::now() + DEADLINE;
    loop {
        let remaining = deadline
            .checked_duration_since(Instant::now())
            .unwrap_or_default();
        assert!(
            !remaining.is_zero(),
            "timed out waiting for condition: {what}"
        );

        // The probe is bounded too. A deadline checked only between probes is
        // no deadline at all: one call into a wedged agent that never answers
        // would hang the test process for ever instead of failing it.
        match tokio::time::timeout(remaining, probe()).await {
            Ok(Some(value)) => return value,
            Ok(None) => tokio::time::sleep(POLL_INTERVAL).await,
            Err(_) => panic!("timed out waiting for condition: {what}"),
        }
    }
}

/// Lets a few discovery rounds pass.
///
/// Only ever used before asserting that something did **not** happen; waiting
/// for success always goes through [`wait_event`] or [`wait_until`].
pub async fn settle() {
    tokio::time::sleep(Duration::from_millis(900)).await;
}

/// Waits until `agent` has `count` authenticated peers in `network`.
pub async fn wait_for_peers(
    agent: &Agent,
    network: crate::NetworkId,
    count: usize,
) -> Vec<iroh::EndpointId> {
    wait_until(&format!("{count} peers in {network}"), || async move {
        let status = agent.network_status(network).await.ok()?;
        if status.peers.len() >= count {
            Some(status.connected_peers())
        } else {
            None
        }
    })
    .await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test(start_paused = true)]
    #[should_panic(expected = "timed out waiting for condition")]
    async fn a_probe_that_never_answers_fails_the_deadline_instead_of_hanging() {
        // Virtual time, so the deadline arrives at once rather than in thirty
        // real seconds — and the deadline does arrive, which is the point: a
        // probe that never returns used to hang the process instead.
        wait_until::<(), _, _>("an answer that never comes", || async {
            std::future::pending::<Option<()>>().await
        })
        .await;
    }
}