dirigent/crates/dirigent_acp_api/tests/common/mod.rs

//! Common test utilities for bidirectional flow testing.
//!
//! This module provides mock implementations and test helpers for testing
//! the bidirectional request/response flow in the ACP Server.

use anyhow::Result;
use serde_json::{json, Value};
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot, Mutex};
use std::collections::HashMap;

/// Mock event for testing event bridge
#[derive(Debug, Clone)]
pub struct MockEvent {
    pub event_type: String,
    pub data: Value,
}

/// Mock SSE client for testing
///
/// Simulates an HTTP client that receives SSE events and posts responses.
pub struct MockSseClient {
    pub client_id: String,
    /// Events received via SSE
    pub received_events: Arc<Mutex<Vec<MockEvent>>>,
    /// Sender for simulating HTTP POST responses
    pub response_tx: mpsc::UnboundedSender<(Value, oneshot::Sender<Result<()>>)>,
}

impl MockSseClient {
    pub fn new(client_id: String) -> (Self, mpsc::UnboundedReceiver<(Value, oneshot::Sender<Result<()>>)>) {
        let (response_tx, response_rx) = mpsc::unbounded_channel();

        (
            Self {
                client_id,
                received_events: Arc::new(Mutex::new(Vec::new())),
                response_tx,
            },
            response_rx,
        )
    }

    /// Simulate receiving an SSE event
    pub async fn receive_sse(&self, event_type: String, data: Value) {
        let mut events = self.received_events.lock().await;
        events.push(MockEvent { event_type, data });
    }

    /// Get all received events
    pub async fn get_events(&self) -> Vec<MockEvent> {
        let events = self.received_events.lock().await;
        events.clone()
    }

    /// Get the most recent event of a specific type
    pub async fn get_latest_event(&self, event_type: &str) -> Option<MockEvent> {
        let events = self.received_events.lock().await;
        events.iter()
            .filter(|e| e.event_type == event_type)
            .last()
            .cloned()
    }

    /// Clear received events
    pub async fn clear_events(&self) {
        let mut events = self.received_events.lock().await;
        events.clear();
    }

    /// Simulate sending a response to /acp/agent_response
    pub async fn send_response(&self, response: Value) -> Result<()> {
        let (tx, rx) = oneshot::channel();
        self.response_tx.send((response, tx))
            .map_err(|_| anyhow::anyhow!("Failed to send response"))?;
        rx.await?
    }
}

/// Mock connector for testing
///
/// Simulates an ACP connector that can send agent requests and receive responses.
pub struct MockConnector {
    pub connector_id: String,
    /// Channel for receiving agent requests from this connector
    pub request_tx: mpsc::UnboundedSender<(Value, String, String, Value)>,
    /// Pending responses (request_id -> sender)
    pub pending_responses: Arc<Mutex<HashMap<Value, oneshot::Sender<Value>>>>,
}

impl MockConnector {
    pub fn new(
        connector_id: String,
    ) -> (Self, mpsc::UnboundedReceiver<(Value, String, String, Value)>) {
        let (request_tx, request_rx) = mpsc::unbounded_channel();

        (
            Self {
                connector_id,
                request_tx,
                pending_responses: Arc::new(Mutex::new(HashMap::new())),
            },
            request_rx,
        )
    }

    /// Simulate sending an agent request
    pub async fn send_agent_request(
        &self,
        session_id: String,
        request_id: Value,
        method: String,
        params: Value,
    ) -> oneshot::Receiver<Value> {
        let (tx, rx) = oneshot::channel();

        // Store the response sender
        let mut pending = self.pending_responses.lock().await;
        pending.insert(request_id.clone(), tx);
        drop(pending);

        // Send the request
        self.request_tx.send((request_id, session_id, method, params))
            .expect("Failed to send agent request");

        rx
    }

    /// Complete a pending response (simulating response from ACP Server)
    pub async fn complete_response(&self, request_id: Value, response: Value) -> Result<()> {
        let mut pending = self.pending_responses.lock().await;

        if let Some(tx) = pending.remove(&request_id) {
            tx.send(response)
                .map_err(|_| anyhow::anyhow!("Failed to send response to connector"))?;
            Ok(())
        } else {
            Err(anyhow::anyhow!("No pending response for request_id: {}", request_id))
        }
    }
}

/// Test context for integration tests
///
/// Provides a complete test environment with mocked components.
pub struct TestContext {
    /// Mock SSE clients by client_id
    pub clients: Arc<Mutex<HashMap<String, MockSseClient>>>,
    /// Mock connectors by connector_id
    pub connectors: Arc<Mutex<HashMap<String, MockConnector>>>,
}

impl TestContext {
    pub fn new() -> Self {
        Self {
            clients: Arc::new(Mutex::new(HashMap::new())),
            connectors: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Create a mock SSE client
    pub async fn create_client(
        &self,
        client_id: String,
    ) -> (MockSseClient, mpsc::UnboundedReceiver<(Value, oneshot::Sender<Result<()>>)>) {
        let (client, response_rx) = MockSseClient::new(client_id.clone());

        let mut clients = self.clients.lock().await;
        clients.insert(client_id.clone(), client.clone());

        (client, response_rx)
    }

    /// Create a mock connector
    pub async fn create_connector(
        &self,
        connector_id: String,
    ) -> (MockConnector, mpsc::UnboundedReceiver<(Value, String, String, Value)>) {
        let (connector, request_rx) = MockConnector::new(connector_id.clone());

        let mut connectors = self.connectors.lock().await;
        connectors.insert(connector_id.clone(), connector.clone());

        (connector, request_rx)
    }

    /// Get a client by ID
    pub async fn get_client(&self, client_id: &str) -> Option<MockSseClient> {
        let clients = self.clients.lock().await;
        clients.get(client_id).cloned()
    }

    /// Get a connector by ID
    pub async fn get_connector(&self, connector_id: &str) -> Option<MockConnector> {
        let connectors = self.connectors.lock().await;
        connectors.get(connector_id).cloned()
    }
}

impl Default for TestContext {
    fn default() -> Self {
        Self::new()
    }
}

/// Helper to create a sample permission request
pub fn sample_permission_request(request_id: u64) -> Value {
    json!({
        "jsonrpc": "2.0",
        "id": request_id,
        "method": "session/request_permission",
        "params": {
            "sessionId": "test-session",
            "tool": "Write",
            "parameters": {
                "path": "/tmp/test.txt",
                "content": "test"
            }
        }
    })
}

/// Helper to create a sample permission response
pub fn sample_permission_response(request_id: u64, allow: bool) -> Value {
    json!({
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "selectedOptionId": if allow { "allow" } else { "deny" }
        }
    })
}

/// Helper to extract agent_request data from SSE event
pub fn extract_agent_request(event: &MockEvent) -> Option<(Value, String, Value)> {
    if event.event_type != "session/update" {
        return None;
    }

    let update = event.data.get("update")?;
    if update.get("sessionUpdate")?.as_str()? != "agent_request" {
        return None;
    }

    let request_id = update.get("requestId")?.clone();
    let method = update.get("method")?.as_str()?.to_string();
    let params = update.get("params")?.clone();

    Some((request_id, method, params))
}

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

    #[tokio::test]
    async fn test_mock_client() {
        let (client, _response_rx) = MockSseClient::new("test-client".to_string());

        // Simulate receiving an event
        client.receive_sse("session/update".to_string(), json!({"test": "data"})).await;

        // Verify event was received
        let events = client.get_events().await;
        assert_eq!(events.len(), 1);
        assert_eq!(events[0].event_type, "session/update");
    }

    #[tokio::test]
    async fn test_mock_connector() {
        let (connector, mut request_rx) = MockConnector::new("test-connector".to_string());

        // Simulate sending an agent request
        let response_fut = connector.send_agent_request(
            "session-1".to_string(),
            json!(0),
            "session/request_permission".to_string(),
            json!({"tool": "Write"}),
        ).await;

        // Verify request was sent
        let request = request_rx.recv().await.unwrap();
        assert_eq!(request.0, json!(0));
        assert_eq!(request.1, "session-1");
        assert_eq!(request.2, "session/request_permission");

        // Simulate completing the response
        connector.complete_response(json!(0), json!({"result": "success"})).await.unwrap();

        // Verify response was received
        let response = response_fut.await.unwrap();
        assert_eq!(response, json!({"result": "success"}));
    }

    #[tokio::test]
    async fn test_test_context() {
        let ctx = TestContext::new();

        // Create client and connector
        let (client, _) = ctx.create_client("client-1".to_string()).await;
        let (connector, _) = ctx.create_connector("connector-1".to_string()).await;

        // Verify we can retrieve them
        assert!(ctx.get_client("client-1").await.is_some());
        assert!(ctx.get_connector("connector-1").await.is_some());
        assert!(ctx.get_client("non-existent").await.is_none());
    }

    #[test]
    fn test_sample_helpers() {
        let request = sample_permission_request(0);
        assert_eq!(request["method"], "session/request_permission");

        let response = sample_permission_response(0, true);
        assert_eq!(response["result"]["selectedOptionId"], "allow");
    }

    #[test]
    fn test_extract_agent_request() {
        let event = MockEvent {
            event_type: "session/update".to_string(),
            data: json!({
                "sessionId": "session-1",
                "update": {
                    "sessionUpdate": "agent_request",
                    "requestId": 0,
                    "method": "session/request_permission",
                    "params": {"tool": "Write"}
                }
            }),
        };

        let (request_id, method, params) = extract_agent_request(&event).unwrap();
        assert_eq!(request_id, json!(0));
        assert_eq!(method, "session/request_permission");
        assert_eq!(params["tool"], "Write");
    }
}