137 lines
4.1 KiB
Rust
137 lines
4.1 KiB
Rust
//! A fake `brokerd` for the `BrokerPort` tests. Included with `#[path]`. Do not edit.
|
|
//!
|
|
//! The fake behaves as the real one will: it reads one request frame, answers on the same
|
|
//! connection, never half-closes, and closes after the final frame.
|
|
|
|
#![allow(dead_code)] // each test file uses a different part of this module
|
|
|
|
use std::os::unix::net::{UnixListener, UnixStream};
|
|
use std::path::PathBuf;
|
|
use std::sync::atomic::{AtomicU32, Ordering};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::thread::{self, JoinHandle};
|
|
use std::time::{Duration, Instant};
|
|
|
|
use loopd::broker_port::{BrokerPort, UNAVAILABLE};
|
|
use loopd::tools::{Pending, ToolPort};
|
|
use proto::{
|
|
CallId, DataClass, Envelope, Message, PROTOCOL_VERSION, SessionId, Timestamp, ToolRequest,
|
|
ToolResponse, read_frame, write_frame,
|
|
};
|
|
|
|
static NEXT: AtomicU32 = AtomicU32::new(0);
|
|
|
|
pub fn socket_path() -> PathBuf {
|
|
let n = NEXT.fetch_add(1, Ordering::SeqCst);
|
|
std::env::temp_dir().join(format!("loopd-broker-{}-{n}.sock", std::process::id()))
|
|
}
|
|
|
|
/// Accepts one connection, reads the request frame and hands both to `script`. The connection
|
|
/// closes when `script` returns.
|
|
pub fn broker<F>(script: F) -> (PathBuf, JoinHandle<Envelope>)
|
|
where
|
|
F: FnOnce(&mut UnixStream, &Envelope) + Send + 'static,
|
|
{
|
|
let path = socket_path();
|
|
let _ = std::fs::remove_file(&path);
|
|
let listener = UnixListener::bind(&path).unwrap();
|
|
let handle = thread::spawn(move || {
|
|
let (mut stream, _) = listener.accept().unwrap();
|
|
let request = read_frame(&mut stream).unwrap();
|
|
script(&mut stream, &request);
|
|
request
|
|
});
|
|
(path, handle)
|
|
}
|
|
|
|
pub fn frame(id: u64, r#final: bool, msg: Message) -> Envelope {
|
|
Envelope {
|
|
v: PROTOCOL_VERSION,
|
|
id,
|
|
r#final,
|
|
msg,
|
|
}
|
|
}
|
|
|
|
pub fn send(stream: &mut UnixStream, id: u64, r#final: bool, response: ToolResponse) {
|
|
// The port may already have given up and gone; the fake does not care.
|
|
let _ = write_frame(stream, &frame(id, r#final, Message::ToolResponse(response)));
|
|
}
|
|
|
|
pub fn result(content: &str) -> ToolResponse {
|
|
ToolResponse::Result {
|
|
content: content.to_string(),
|
|
class: DataClass::Secret,
|
|
untrusted: true,
|
|
truncated: false,
|
|
}
|
|
}
|
|
|
|
pub fn request() -> ToolRequest {
|
|
ToolRequest {
|
|
session: SessionId::new("chat-1").unwrap(),
|
|
call: CallId(7),
|
|
tool: "read_file".to_string(),
|
|
arguments: r#"{"path":"/home/kyle/notes/a.md"}"#.to_string(),
|
|
}
|
|
}
|
|
|
|
pub fn in_ms(ms: u64) -> Timestamp {
|
|
Timestamp::from_unix_millis(Timestamp::now().unix_millis() + ms).unwrap()
|
|
}
|
|
|
|
pub struct Call {
|
|
pub response: ToolResponse,
|
|
pub pending: Vec<Pending>,
|
|
pub lines: Vec<String>,
|
|
pub took: Duration,
|
|
}
|
|
|
|
pub fn call(socket: PathBuf, timeout_ms: u64, request: &ToolRequest) -> Call {
|
|
let lines = Arc::new(Mutex::new(Vec::new()));
|
|
let sink = lines.clone();
|
|
let port = BrokerPort::with_log(
|
|
socket,
|
|
Duration::from_millis(timeout_ms),
|
|
Box::new(move |line| sink.lock().unwrap().push(line.to_string())),
|
|
);
|
|
let mut pending = Vec::new();
|
|
let started = Instant::now();
|
|
let response = port.call(request, &mut |p| pending.push(*p));
|
|
let took = started.elapsed();
|
|
let lines = lines.lock().unwrap().clone();
|
|
Call {
|
|
response,
|
|
pending,
|
|
lines,
|
|
took,
|
|
}
|
|
}
|
|
|
|
/// The answer is the fixed failure, and `loopd` printed exactly one line, with the pointer.
|
|
pub fn assert_unavailable(call: &Call, why: &str) {
|
|
assert_eq!(
|
|
call.response,
|
|
ToolResponse::Failed {
|
|
message: UNAVAILABLE.to_string()
|
|
},
|
|
"{why}"
|
|
);
|
|
assert_eq!(
|
|
call.lines.len(),
|
|
1,
|
|
"{why}: one line per failed call: {:?}",
|
|
call.lines
|
|
);
|
|
let line = &call.lines[0];
|
|
assert!(
|
|
line.starts_with("loopd: the tool broker is unavailable: "),
|
|
"{why}: {line}"
|
|
);
|
|
assert!(
|
|
line.ends_with("; see docs/runbook.md#broker-unavailable"),
|
|
"{why}: {line}"
|
|
);
|
|
assert!(call.pending.is_empty() || why.contains("pending"), "{why}");
|
|
}
|