use std::{
net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener},
path::PathBuf,
process::{Command, Output, Stdio},
time::Duration,
};
use datum::{Keep, Sink, Source};
use datum_agent::{
Agent, AgentConfig, JobMat, JobSpec,
dcp::{DcpJobFactories, DcpServer, DcpServerConfig, DcpServerHandle, DcpTcpServerConfig},
};
use serde_json::Value;
use tokio::{
io::{AsyncBufReadExt, BufReader},
process::Command as TokioCommand,
};
struct TestDaemon {
agent: datum_agent::AgentHandle,
handle: DcpServerHandle,
addr: SocketAddr,
}
impl TestDaemon {
async fn start() -> Self {
let agent = Agent::start_with_config(AgentConfig {
poll_interval: Duration::from_millis(5),
event_buffer: 128,
..AgentConfig::default()
})
.expect("agent starts");
let server = DcpServer::from_agent(
&agent,
factories(),
DcpServerConfig {
node_id: "cli-test-agent".to_owned(),
tcp: Some(DcpTcpServerConfig {
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
}),
quic: None,
auth_token: None,
metrics_interval: Duration::from_millis(10),
frame_buffer: 128,
},
);
let handle = server.start().await.expect("DCP server starts");
let addr = handle.tcp_addr().expect("TCP listener address");
Self {
agent,
handle,
addr,
}
}
async fn shutdown(self) {
self.handle.shutdown().await;
self.agent
.registry()
.shutdown()
.expect("registry shuts down");
}
}
fn factories() -> DcpJobFactories {
let factories = DcpJobFactories::new();
factories
.register("ticker", |instance_name, _params| {
Ok(JobSpec::new(instance_name, |context| {
let control = context.control();
Ok(
Source::tick(Duration::ZERO, Duration::from_millis(10), 1_u64)
.instrumented(
format!("{}:{}", context.name(), context.generation()),
context.instrumentation_registry(),
)
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}),
)
}))
})
.expect("ticker factory registers");
factories
.register("empty", |instance_name, _params| {
Ok(JobSpec::new(instance_name, |context| {
let control = context.control();
Ok(Source::<u64>::empty()
.instrumented(
format!("{}:{}", context.name(), context.generation()),
context.instrumentation_registry(),
)
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}))
}))
})
.expect("empty factory registers");
factories
}
fn datum_bin() -> PathBuf {
PathBuf::from(env!("CARGO_BIN_EXE_datum"))
}
fn run_datum(addr: SocketAddr, args: &[&str]) -> Output {
Command::new(datum_bin())
.arg("--addr")
.arg(addr.to_string())
.args(args)
.output()
.expect("datum process runs")
}
fn assert_success(output: &Output) {
assert!(
output.status.success(),
"expected success\nstdout:\n{}\nstderr:\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
}
fn stdout(output: &Output) -> String {
String::from_utf8(output.stdout.clone()).expect("stdout is UTF-8")
}
fn stderr(output: &Output) -> String {
String::from_utf8(output.stderr.clone()).expect("stderr is UTF-8")
}
async fn stop_child(child: &mut tokio::process::Child) {
let _ = child.start_kill();
let _ = tokio::time::timeout(Duration::from_secs(1), child.wait()).await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn lifecycle_commands_and_json_snapshots() {
let daemon = TestDaemon::start().await;
let output = run_datum(daemon.addr, &["--json", "ps"]);
assert_success(&output);
assert_eq!(stdout(&output), "{\n \"jobs\": []\n}\n");
let output = run_datum(
daemon.addr,
&["start", "ticker", "--name", "ingest", "--param", "shard=0"],
);
assert_success(&output);
assert!(stdout(&output).contains("started 'ingest' from factory 'ticker'"));
let output = run_datum(daemon.addr, &["ps"]);
assert_success(&output);
let jobs_table = stdout(&output);
assert!(jobs_table.contains("JOB"));
assert!(jobs_table.contains("ingest"));
assert!(jobs_table.contains("Running"));
let output = run_datum(daemon.addr, &["--json", "status", "ingest"]);
assert_success(&output);
let status: Value = serde_json::from_slice(&output.stdout).expect("status JSON parses");
assert_eq!(status["job"]["name"], "ingest");
assert_eq!(status["job"]["state"], "Running");
assert_eq!(status["job"]["generation"], 1);
let output = run_datum(daemon.addr, &["--json", "restart", "ingest"]);
assert_success(&output);
let restarted: Value = serde_json::from_slice(&output.stdout).expect("restart JSON parses");
assert_eq!(restarted["action"], "restarted");
assert_eq!(restarted["job"]["name"], "ingest");
assert_eq!(restarted["job"]["generation"], 2);
let output = run_datum(daemon.addr, &["--json", "metrics", "ingest"]);
assert_success(&output);
let metrics: Value = serde_json::from_slice(&output.stdout).expect("metrics JSON parses");
assert_eq!(metrics["job"], "ingest");
assert!(
metrics["sample"]["streams"]
.as_array()
.expect("metric streams array")
.iter()
.any(|stream| stream["name"]
.as_str()
.is_some_and(|name| name.starts_with("ingest:")))
);
let output = run_datum(daemon.addr, &["drain", "ingest"]);
assert_success(&output);
assert!(stdout(&output).contains("drained 'ingest'"));
let output = run_datum(daemon.addr, &["start", "ticker", "--name", "stop-me"]);
assert_success(&output);
let output = run_datum(daemon.addr, &["--json", "stop", "stop-me"]);
assert_success(&output);
let stopped: Value = serde_json::from_slice(&output.stdout).expect("stop JSON parses");
assert_eq!(stopped["action"], "stopped");
assert_eq!(stopped["job"]["state"], "Stopped");
daemon.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn events_command_tails_subscription() {
let daemon = TestDaemon::start().await;
let child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(daemon.addr.to_string())
.arg("events")
.stdout(Stdio::piped())
.spawn()
.expect("events command starts");
tokio::time::sleep(Duration::from_millis(100)).await;
let output = run_datum(daemon.addr, &["start", "ticker", "--name", "event-job"]);
assert_success(&output);
let output = tokio::time::timeout(Duration::from_secs(4), child.wait_with_output())
.await
.expect("events command exits after idle")
.expect("events output");
assert_success(&output);
let output = stdout(&output);
assert!(output.contains("event-job"));
assert!(output.contains("Submitted") || output.contains("Started"));
daemon.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn follow_streams_emit_json_lines() {
let daemon = TestDaemon::start().await;
let mut child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(daemon.addr.to_string())
.arg("--json")
.arg("events")
.arg("--follow")
.stdout(Stdio::piped())
.spawn()
.expect("events --follow command starts");
tokio::time::sleep(Duration::from_millis(100)).await;
let output = run_datum(daemon.addr, &["start", "ticker", "--name", "follow-events"]);
assert_success(&output);
let stdout = child.stdout.take().expect("event stdout is piped");
let mut reader = BufReader::new(stdout);
let mut line = String::new();
tokio::time::timeout(Duration::from_secs(4), reader.read_line(&mut line))
.await
.expect("event line arrives")
.expect("event line reads");
drop(reader);
stop_child(&mut child).await;
let event: Value = serde_json::from_str(&line).expect("event JSON line parses");
assert_eq!(event["event"]["name"], "follow-events");
let output = run_datum(
daemon.addr,
&["start", "ticker", "--name", "follow-metrics"],
);
assert_success(&output);
let mut child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(daemon.addr.to_string())
.arg("--json")
.arg("metrics")
.arg("follow-metrics")
.arg("--follow")
.stdout(Stdio::piped())
.spawn()
.expect("metrics --follow command starts");
let stdout = child.stdout.take().expect("metric stdout is piped");
let mut reader = BufReader::new(stdout);
let mut line = String::new();
tokio::time::timeout(Duration::from_secs(4), reader.read_line(&mut line))
.await
.expect("metric line arrives")
.expect("metric line reads");
drop(reader);
stop_child(&mut child).await;
let metric: Value = serde_json::from_str(&line).expect("metric JSON line parses");
assert_eq!(metric["job"], "follow-metrics");
assert!(
metric["sample"]["streams"]
.as_array()
.expect("metric streams array")
.iter()
.any(|stream| stream["name"]
.as_str()
.is_some_and(|name| name.starts_with("follow-metrics:")))
);
daemon.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn connect_failure_uses_exit_code_2_and_fix_text() {
let listener = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0))
.expect("ephemeral port binds");
let addr = listener.local_addr().expect("listener address");
drop(listener);
let output = run_datum(addr, &["ps"]);
assert_eq!(output.status.code(), Some(2));
let stderr = stderr(&output);
assert!(stderr.contains("could not connect to datum-agent"));
assert!(stderr.contains("start datum-agent or pass --addr"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_errors_use_exit_code_1_and_fix_text() {
let daemon = TestDaemon::start().await;
let output = run_datum(daemon.addr, &["status", "missing"]);
assert_eq!(output.status.code(), Some(1));
let stderr = stderr(&output);
assert!(stderr.contains("datum-agent rejected the request"));
assert!(stderr.contains("Check the job or factory name with `datum ps`."));
daemon.shutdown().await;
}
#[test]
fn usage_errors_use_exit_code_2() {
let output = Command::new(datum_bin())
.arg("start")
.arg("ticker")
.arg("--param")
.arg("bad-param")
.output()
.expect("datum process runs");
assert_eq!(output.status.code(), Some(2));
assert!(stderr(&output).contains("invalid --param 'bad-param'"));
}