maestro/crates/maestro-testing/src/lib.rs

//! An in-memory Maestro daemon for tests.
//!
//! Serves the real `/api/*` surface from `maestro-protocol` types, backed by a
//! `Vec<Session>` instead of tmux. Use it whenever the subject under test is a
//! client — request building, error decoding, the catalog reducer, the proxy —
//! rather than session management itself. Tests that genuinely need tmux belong
//! in `maestro-e2e` against a real daemon.

#![deny(missing_debug_implementations)]

use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr};
use std::sync::{Arc, Mutex};

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Path, Query, State};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::{any, delete, get};
use axum::{Json, Router};
use chrono::Utc;
use futures_util::{SinkExt, StreamExt};
use maestro_protocol::{
    routes, ApiError, CatalogEvent, CreateSessionRequest, DaemonClient, DaemonInfo, DirListing,
    GitRoot, RecentDirs, Session, SessionStatus, PROTOCOL_VERSION,
};
use tokio::net::TcpListener;
use tokio::sync::{broadcast, oneshot};
use uuid::Uuid;

/// Build a plausible `Session` for a test to hand to [`MockDaemon::add_session`].
pub fn session(id: Uuid, cwd: &str, command: &str) -> Session {
    Session {
        id,
        cwd: cwd.into(),
        command: command.into(),
        env: HashMap::new(),
        metadata: HashMap::new(),
        worktree: None,
        status: SessionStatus::Running,
        created_at: Utc::now(),
        foreground: None,
        activity: Default::default(),
        last_activity_at: Utc::now(),
        activity_since: chrono::Utc::now(),
        current_cwd: None,
        blocked_work: None,
    }
}

#[derive(Debug, Default)]
struct Inner {
    sessions: Vec<Session>,
    info: Option<DaemonInfo>,
    /// When set, every route answers with this instead of its normal response.
    forced_error: Option<(u16, ApiError)>,
}

#[derive(Clone)]
struct AppState {
    inner: Arc<Mutex<Inner>>,
    events: broadcast::Sender<CatalogEvent>,
}

impl std::fmt::Debug for AppState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AppState").finish_non_exhaustive()
    }
}

/// A running mock daemon. Shuts down when dropped.
#[derive(Debug)]
pub struct MockDaemon {
    addr: SocketAddr,
    state: AppState,
    _shutdown: oneshot::Sender<()>,
}

impl MockDaemon {
    pub async fn start() -> Self {
        let state = AppState {
            inner: Arc::new(Mutex::new(Inner::default())),
            events: broadcast::channel(64).0,
        };
        let app = Router::new()
            .route(routes::SESSIONS, get(list_sessions).post(create_session))
            .route(routes::patterns::SESSION, delete(kill_session))
            .route(routes::INFO, get(daemon_info))
            .route(routes::DIRS, get(list_dir))
            .route(routes::RECENT_DIRS, get(recent_dirs))
            .route(routes::GITROOT, get(git_root))
            .route(routes::EVENTS, any(events_ws))
            .route(routes::patterns::ATTACH, any(attach_ws))
            .with_state(state.clone());

        let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
            .await
            .expect("bind");
        let addr = listener.local_addr().expect("local_addr");
        let (tx, rx) = oneshot::channel::<()>();
        tokio::spawn(async move {
            let _ = axum::serve(listener, app)
                .with_graceful_shutdown(async move {
                    let _ = rx.await;
                })
                .await;
        });

        Self {
            addr,
            state,
            _shutdown: tx,
        }
    }

    pub fn addr(&self) -> SocketAddr {
        self.addr
    }

    pub fn http_base(&self) -> String {
        format!("http://{}", self.addr)
    }

    /// A client bound to this daemon.
    pub fn client(&self) -> DaemonClient {
        DaemonClient::new(self.http_base())
    }

    /// Add a session and announce it, as a real spawn would.
    pub fn add_session(&self, session: Session) {
        self.state
            .inner
            .lock()
            .expect("poisoned")
            .sessions
            .push(session.clone());
        let _ = self
            .state
            .events
            .send(CatalogEvent::SessionAdded { session });
    }

    /// Remove a session and announce it.
    pub fn remove_session(&self, id: Uuid) {
        self.state
            .inner
            .lock()
            .expect("poisoned")
            .sessions
            .retain(|s| s.id != id);
        let _ = self.state.events.send(CatalogEvent::SessionRemoved { id });
    }

    pub fn set_info(&self, info: DaemonInfo) {
        self.state.inner.lock().expect("poisoned").info = Some(info);
    }

    /// Make every route answer with this error, so clients' failure paths — and
    /// their [`ApiError`] decoding — can be exercised.
    pub fn force_error(&self, status: u16, code: &str, message: &str) {
        self.state.inner.lock().expect("poisoned").forced_error =
            Some((status, ApiError::new(code, message)));
    }

    pub fn clear_error(&self) {
        self.state.inner.lock().expect("poisoned").forced_error = None;
    }

    pub fn session_count(&self) -> usize {
        self.state.inner.lock().expect("poisoned").sessions.len()
    }
}

fn forced(state: &AppState) -> Option<Response> {
    let guard = state.inner.lock().expect("poisoned");
    guard.forced_error.as_ref().map(|(status, body)| {
        (
            StatusCode::from_u16(*status).expect("valid status"),
            Json(body.clone()),
        )
            .into_response()
    })
}

async fn list_sessions(State(s): State<AppState>) -> Response {
    if let Some(r) = forced(&s) {
        return r;
    }
    let sessions = s.inner.lock().expect("poisoned").sessions.clone();
    Json(sessions).into_response()
}

async fn create_session(
    State(s): State<AppState>,
    Json(req): Json<CreateSessionRequest>,
) -> Response {
    if let Some(r) = forced(&s) {
        return r;
    }
    let mut created = session(Uuid::new_v4(), &req.cwd.to_string_lossy(), &req.command);
    created.env = req.env;
    created.metadata = req.metadata;
    s.inner
        .lock()
        .expect("poisoned")
        .sessions
        .push(created.clone());
    let _ = s.events.send(CatalogEvent::SessionAdded {
        session: created.clone(),
    });
    Json(created).into_response()
}

async fn kill_session(State(s): State<AppState>, Path(id): Path<Uuid>) -> Response {
    if let Some(r) = forced(&s) {
        return r;
    }
    let existed = {
        let mut guard = s.inner.lock().expect("poisoned");
        let before = guard.sessions.len();
        guard.sessions.retain(|x| x.id != id);
        guard.sessions.len() != before
    };
    if !existed {
        return (
            StatusCode::NOT_FOUND,
            Json(ApiError::new("not_found", "no such session on this daemon")),
        )
            .into_response();
    }
    let _ = s.events.send(CatalogEvent::SessionRemoved { id });
    StatusCode::NO_CONTENT.into_response()
}

async fn daemon_info(State(s): State<AppState>) -> Response {
    if let Some(r) = forced(&s) {
        return r;
    }
    let info = s
        .inner
        .lock()
        .expect("poisoned")
        .info
        .clone()
        .unwrap_or_else(|| DaemonInfo {
            hostname: "mock-host".into(),
            version: PROTOCOL_VERSION.into(),
            home_dir: "/home/mock".into(),
            presets: Vec::new(),
        });
    Json(info).into_response()
}

async fn list_dir(State(s): State<AppState>, Query(q): Query<PathQuery>) -> Response {
    if let Some(r) = forced(&s) {
        return r;
    }
    Json(DirListing {
        path: q.path.into(),
        entries: Vec::new(),
        truncated: false,
    })
    .into_response()
}

async fn recent_dirs(State(s): State<AppState>) -> Response {
    if let Some(r) = forced(&s) {
        return r;
    }
    Json(RecentDirs {
        dirs: Vec::new(),
        truncated: false,
    })
    .into_response()
}

async fn git_root(State(s): State<AppState>, Query(_q): Query<PathQuery>) -> Response {
    if let Some(r) = forced(&s) {
        return r;
    }
    Json(GitRoot { root: None }).into_response()
}

/// Snapshot on connect, then live deltas — the same contract the daemon offers.
async fn events_ws(State(s): State<AppState>, ws: WebSocketUpgrade) -> Response {
    ws.on_upgrade(move |socket| events_loop(s, socket))
}

async fn events_loop(s: AppState, socket: WebSocket) {
    let (mut sink, mut stream) = socket.split();
    let mut rx = s.events.subscribe();
    let snapshot = CatalogEvent::Snapshot {
        sessions: s.inner.lock().expect("poisoned").sessions.clone(),
    };
    if sink
        .send(Message::Text(
            serde_json::to_string(&snapshot).expect("encode"),
        ))
        .await
        .is_err()
    {
        return;
    }
    loop {
        tokio::select! {
            ev = rx.recv() => {
                let Ok(ev) = ev else { break };
                let text = serde_json::to_string(&ev).expect("encode");
                if sink.send(Message::Text(text)).await.is_err() {
                    break;
                }
            }
            frame = stream.next() => {
                match frame {
                    None | Some(Err(_)) | Some(Ok(Message::Close(_))) => break,
                    Some(Ok(_)) => {}
                }
            }
        }
    }
}

/// Echoes whatever it is sent, so a client's PTY plumbing can be exercised.
async fn attach_ws(
    State(s): State<AppState>,
    Path(id): Path<Uuid>,
    ws: WebSocketUpgrade,
) -> Response {
    let known = s
        .inner
        .lock()
        .expect("poisoned")
        .sessions
        .iter()
        .any(|x| x.id == id);
    if !known {
        return (
            StatusCode::NOT_FOUND,
            Json(ApiError::new("not_found", "no such session on this daemon")),
        )
            .into_response();
    }
    ws.on_upgrade(|mut socket| async move {
        while let Some(Ok(msg)) = socket.next().await {
            match msg {
                Message::Binary(b) => {
                    if socket.send(Message::Binary(b)).await.is_err() {
                        break;
                    }
                }
                Message::Text(t) => {
                    if socket.send(Message::Text(t)).await.is_err() {
                        break;
                    }
                }
                Message::Close(_) => break,
                _ => {}
            }
        }
    })
}

#[derive(Debug, serde::Deserialize)]
struct PathQuery {
    #[serde(default)]
    path: String,
}