use crate::agent::files::{load_json, AgentFile, AGENT_FILE};
use std::path::Path;
use std::time::Duration;
pub struct AgentClient {
base: String,
port: u16,
token: String,
}
impl AgentClient {
pub fn from_dir(agent_dir: &Path) -> Option<Self> {
let f: AgentFile = load_json(&agent_dir.join(AGENT_FILE))?;
Some(Self {
base: format!("http://127.0.0.1:{}", f.port),
port: f.port,
token: f.token,
})
}
pub fn port(&self) -> u16 {
self.port
}
pub fn running() -> Option<Self> {
let c = Self::from_dir(&crate::credentials::dir()?.join("agent"))?;
c.summary(Duration::from_secs(1)).ok()?;
Some(c)
}
fn http(timeout: Duration) -> ureq::Agent {
ureq::Agent::new_with_config(
ureq::Agent::config_builder()
.http_status_as_error(false)
.timeout_global(Some(timeout))
.build(),
)
}
fn read(resp: ureq::http::Response<ureq::Body>) -> Result<serde_json::Value, String> {
let status = resp.status().as_u16();
let body: serde_json::Value = resp.into_body().read_json().unwrap_or_default();
if status == 200 {
Ok(body)
} else {
Err(body["error"]["message"]
.as_str()
.map(str::to_string)
.unwrap_or_else(|| format!("agent answered HTTP {status}")))
}
}
pub fn summary(&self, timeout: Duration) -> Result<serde_json::Value, String> {
let resp = Self::http(timeout)
.get(&format!("{}/v1/summary", self.base))
.header("Authorization", format!("Bearer {}", self.token).as_str())
.call()
.map_err(|e| e.to_string())?;
Self::read(resp)
}
pub fn put(
&self,
path: &str,
body: serde_json::Value,
timeout: Duration,
) -> Result<serde_json::Value, String> {
let resp = Self::http(timeout)
.put(&format!("{}{path}", self.base))
.header("Authorization", format!("Bearer {}", self.token).as_str())
.send_json(body)
.map_err(|e| e.to_string())?;
Self::read(resp)
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use crate::agent::files::tests::tmp;
pub(crate) struct AgentGuard {
stop: Option<std::sync::mpsc::Sender<()>>,
thread: Option<std::thread::JoinHandle<()>>,
}
impl AgentGuard {
pub(crate) fn new(
stop: std::sync::mpsc::Sender<()>,
thread: std::thread::JoinHandle<()>,
) -> Self {
Self {
stop: Some(stop),
thread: Some(thread),
}
}
}
impl Drop for AgentGuard {
fn drop(&mut self) {
if let Some(stop) = self.stop.take() {
let _ = stop.send(());
}
if let Some(thread) = self.thread.take() {
let _ = thread.join();
}
}
}
pub(crate) fn sleeper_cfg(tag: &str) -> crate::agent::config::AgentConfig {
let free = || {
std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port()
};
let sleeper = vec![
"sh".to_string(),
"-c".to_string(),
"exec sleep 30".to_string(),
];
let mut cfg = crate::agent::config::AgentConfig::for_dirs(tmp(tag));
cfg.port = 0;
cfg.broker_port = free();
cfg.max_workers = 3;
cfg.broker_argv = sleeper.clone();
cfg.worker_template = sleeper;
cfg.worker_template_overridden = true;
cfg.timings.shutdown_wait = Duration::from_secs(2);
cfg
}
pub(crate) fn start_agent(cfg: crate::agent::config::AgentConfig) -> (AgentClient, AgentGuard) {
let dir = cfg.dir.clone();
let (_, _, stop, thread) = crate::agent::run::start_background(cfg).expect("agent starts");
(
AgentClient::from_dir(&dir).expect("agent.json written"),
AgentGuard::new(stop, thread),
)
}
pub(crate) fn idle_agent() -> (AgentClient, AgentGuard) {
start_agent(sleeper_cfg("client"))
}
#[test]
fn talks_to_a_running_agent_and_surfaces_api_errors() {
let (c, _guard) = idle_agent();
let s = c.summary(Duration::from_secs(2)).unwrap();
assert_eq!(s["v"], 1);
let s = c
.put(
"/v1/workers",
serde_json::json!({"count": 2}),
Duration::from_secs(2),
)
.unwrap();
assert_eq!(s["this_mac"]["workers"]["desired"], 2);
let err = c
.put(
"/v1/workers",
serde_json::json!({"count": 9}),
Duration::from_secs(2),
)
.unwrap_err();
assert!(err.contains("count must be between 0 and 3"), "{err}");
}
#[test]
fn a_stale_agent_file_is_not_a_running_agent() {
let dir = tmp("stale");
crate::agent::files::ensure_dir(&dir).unwrap();
crate::agent::files::save_json(
&dir.join(crate::agent::files::AGENT_FILE),
&crate::agent::files::AgentFile {
v: 1,
port: 1,
token: "t".into(),
pid: 1,
version: "x".into(),
},
)
.unwrap();
let c = AgentClient::from_dir(&dir).unwrap();
assert!(c.summary(Duration::from_secs(1)).is_err());
assert!(AgentClient::from_dir(&tmp("absent")).is_none());
}
}