sync from monorepo @ 2452e92e
This commit is contained in:
@@ -0,0 +1,338 @@
|
||||
//! 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");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user