maestro/crates/maestro-e2e/src/client.rs

//! Thin test harness over [`maestro_protocol::DaemonClient`].
//!
//! HTTP calls delegate to the shared client so the harness exercises the same
//! request building, timeouts and error decoding the CLI and TUI use. Only the
//! WebSocket helpers are test-specific: they hand back a raw stream, or a
//! handle that accumulates events in the background.

#![allow(clippy::result_large_err)]

use std::sync::Arc;
use std::time::Duration;

use futures_util::{SinkExt, StreamExt};
use maestro_protocol::{
    BranchInfo, CatalogEvent, CreateSessionRequest, DaemonClient, Project, Session,
};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message as WsMessage;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
use uuid::Uuid;

#[derive(Debug, thiserror::Error)]
#[allow(clippy::result_large_err)]
pub enum ClientError {
    #[error(transparent)]
    Api(#[from] maestro_protocol::ClientError),
    #[error("ws: {0}")]
    Ws(#[from] tokio_tungstenite::tungstenite::Error),
    #[error("json: {0}")]
    Json(#[from] serde_json::Error),
}

/// Daemon client plus a few WS-aware helpers, bound to one daemon.
#[derive(Debug, Clone)]
pub struct TestClient {
    daemon: DaemonClient,
}

/// Per-request budget. Well below the client default: a hung daemon should fail
/// the test, not stall the suite for 30s.
const TEST_TIMEOUT: Duration = Duration::from_secs(5);

impl TestClient {
    pub fn new(http_base: &str) -> Self {
        Self {
            daemon: DaemonClient::with_timeouts(http_base, TEST_TIMEOUT, TEST_TIMEOUT),
        }
    }

    /// The shared client underneath, for calls this harness doesn't wrap.
    pub fn daemon(&self) -> &DaemonClient {
        &self.daemon
    }

    pub async fn list_sessions(&self) -> Result<Vec<Session>, ClientError> {
        Ok(self.daemon.list_sessions().await?)
    }

    pub async fn spawn_session(&self, req: &CreateSessionRequest) -> Result<Session, ClientError> {
        Ok(self.daemon.create_session(req).await?)
    }

    pub async fn kill_session(&self, id: Uuid) -> Result<(), ClientError> {
        Ok(self.daemon.kill_session(id, false).await.map(|_| ())?)
    }

    /// Kill with `force=true`: overrides a kill refusal and leaves the
    /// guarded worktree on disk (#81).
    pub async fn kill_session_forced(&self, id: Uuid) -> Result<(), ClientError> {
        Ok(self.daemon.kill_session(id, true).await?)
    }

    pub async fn list_projects(&self) -> Result<Vec<Project>, ClientError> {
        Ok(self.daemon.list_projects().await?)
    }

    pub async fn list_branches(&self, path: &str) -> Result<Vec<BranchInfo>, ClientError> {
        Ok(self.daemon.list_branches(path).await?)
    }

    /// Connect to the daemon's per-session WS and return the stream. The
    /// caller drives the conversation.
    pub async fn attach(
        &self,
        id: Uuid,
    ) -> Result<WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>, ClientError> {
        let request = self.daemon.attach_url(id)?.into_client_request()?;
        let (ws, _resp) = connect_async(request).await?;
        Ok(ws)
    }

    /// Subscribe to the daemon's `/api/events` WS. Returns a handle whose
    /// background task pushes parsed events into a shared vector.
    pub async fn subscribe_events(&self) -> EventStream {
        let url = self.daemon.events_url().expect("events url");
        let request = url.into_client_request().expect("ws req");
        let (ws, _resp) = connect_async(request).await.expect("events ws");

        let events = Arc::new(tokio::sync::Mutex::new(Vec::new()));
        let stream_events = events.clone();
        let (mut sink, mut stream) = ws.split();
        let task = tokio::spawn(async move {
            while let Some(msg) = stream.next().await {
                match msg {
                    Ok(WsMessage::Text(t)) => {
                        if let Ok(ev) = serde_json::from_str::<CatalogEvent>(&t) {
                            stream_events.lock().await.push(ev);
                        }
                    }
                    Ok(WsMessage::Close(_)) | Err(_) => break,
                    _ => {}
                }
            }
            let _ = sink.close().await;
        });
        EventStream { events, task }
    }
}

/// Handle to a background task draining the events WS into a vector.
#[derive(Debug)]
pub struct EventStream {
    events: Arc<tokio::sync::Mutex<Vec<CatalogEvent>>>,
    task: tokio::task::JoinHandle<()>,
}

impl EventStream {
    pub async fn snapshot(&self) -> Vec<CatalogEvent> {
        self.events.lock().await.clone()
    }
}

impl Drop for EventStream {
    fn drop(&mut self) {
        self.task.abort();
    }
}