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

//! Thin HTTP+WS client used by tests to drive a daemon's `/api/*` surface.

#![allow(clippy::result_large_err)]

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

use futures_util::{SinkExt, StreamExt};
use maestro_protocol::{CatalogEvent, CreateSessionRequest, Session};
use reqwest::Client as HttpClient;
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("http: {0}")]
    Http(#[from] reqwest::Error),
    #[error("ws: {0}")]
    Ws(#[from] tokio_tungstenite::tungstenite::Error),
    #[error("json: {0}")]
    Json(#[from] serde_json::Error),
    #[error("unexpected status {status}: {body}")]
    Status { status: u16, body: String },
}

/// HTTP client plus a few WS-aware helpers, bound to one daemon.
#[derive(Debug, Clone)]
pub struct TestClient {
    http_base: String,
    ws_base: String,
    http: HttpClient,
}

impl TestClient {
    pub fn new(http_base: &str, ws_base: &str) -> Self {
        let http = HttpClient::builder()
            .timeout(Duration::from_secs(5))
            .build()
            .expect("reqwest client");
        Self {
            http_base: http_base.trim_end_matches('/').to_string(),
            ws_base: ws_base.trim_end_matches('/').to_string(),
            http,
        }
    }

    pub async fn list_sessions(&self) -> Result<Vec<Session>, ClientError> {
        let r = self
            .http
            .get(format!("{}/api/sessions", self.http_base))
            .send()
            .await?;
        check_ok(&r)?;
        Ok(r.json().await?)
    }

    pub async fn spawn_session(&self, req: &CreateSessionRequest) -> Result<Session, ClientError> {
        let r = self
            .http
            .post(format!("{}/api/sessions", self.http_base))
            .json(req)
            .send()
            .await?;
        let status = r.status();
        if !status.is_success() {
            let body = r.text().await.unwrap_or_default();
            return Err(ClientError::Status {
                status: status.as_u16(),
                body,
            });
        }
        Ok(r.json().await?)
    }

    pub async fn kill_session(&self, id: Uuid) -> Result<(), ClientError> {
        let r = self
            .http
            .delete(format!("{}/api/sessions/{id}", self.http_base))
            .send()
            .await?;
        check_ok(&r)?;
        Ok(())
    }

    /// 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 url = format!("{}/api/attach/{id}", self.ws_base);
        let request = url.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 = format!("{}/api/events", self.ws_base);
        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 }
    }
}

fn check_ok(r: &reqwest::Response) -> Result<(), ClientError> {
    let status = r.status();
    if !status.is_success() {
        return Err(ClientError::Status {
            status: status.as_u16(),
            body: format!("HTTP {status}"),
        });
    }
    Ok(())
}

/// 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();
    }
}