feat: add first page with auth and containers list and agents
This commit is contained in:
385
agent/src/docker.rs
Normal file
385
agent/src/docker.rs
Normal file
@ -0,0 +1,385 @@
|
||||
use anyhow::Result;
|
||||
use bollard::{
|
||||
container::{
|
||||
ListContainersOptions, LogOutput, LogsOptions, RemoveContainerOptions,
|
||||
StartContainerOptions, StopContainerOptions,
|
||||
},
|
||||
Docker,
|
||||
};
|
||||
use futures_util::Stream;
|
||||
use std::{collections::HashMap, pin::Pin};
|
||||
|
||||
use crate::proto::{ContainerInfo, ContainerPort};
|
||||
|
||||
// ── Public trait ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Abstraction over Docker operations, allowing tests to provide a mock backend.
|
||||
pub trait ContainerBackend: Clone + Send + Sync + 'static {
|
||||
fn list_containers(
|
||||
&self,
|
||||
) -> impl std::future::Future<Output = Result<Vec<ContainerInfo>>> + Send;
|
||||
|
||||
fn start(&self, id: &str) -> impl std::future::Future<Output = Result<()>> + Send;
|
||||
fn stop(&self, id: &str) -> impl std::future::Future<Output = Result<()>> + Send;
|
||||
fn restart(&self, id: &str) -> impl std::future::Future<Output = Result<()>> + Send;
|
||||
fn remove(&self, id: &str) -> impl std::future::Future<Output = Result<()>> + Send;
|
||||
|
||||
fn logs(
|
||||
&self,
|
||||
id: &str,
|
||||
follow: bool,
|
||||
tail: i32,
|
||||
) -> Pin<Box<dyn Stream<Item = Result<LogOutput, bollard::errors::Error>> + Send>>;
|
||||
}
|
||||
|
||||
// ── Real implementation ───────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct DockerClient {
|
||||
inner: Docker,
|
||||
}
|
||||
|
||||
impl DockerClient {
|
||||
pub fn new() -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: Docker::connect_with_socket_defaults()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ContainerBackend for DockerClient {
|
||||
async fn list_containers(&self) -> Result<Vec<ContainerInfo>> {
|
||||
let opts = ListContainersOptions::<String> {
|
||||
all: true,
|
||||
..Default::default()
|
||||
};
|
||||
let containers = self.inner.list_containers(Some(opts)).await?;
|
||||
|
||||
let result = containers
|
||||
.into_iter()
|
||||
.map(|c| {
|
||||
let id = c.id.unwrap_or_default();
|
||||
let name = c
|
||||
.names
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim_start_matches('/')
|
||||
.to_string();
|
||||
|
||||
let ports = c
|
||||
.ports
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|p| ContainerPort {
|
||||
host_port: p.public_port.unwrap_or(0) as u32,
|
||||
container_port: p.private_port as u32,
|
||||
protocol: p
|
||||
.typ
|
||||
.map(|t| format!("{:?}", t).to_lowercase())
|
||||
.unwrap_or_default(),
|
||||
host_ip: p.ip.unwrap_or_default(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let labels: HashMap<String, String> = c.labels.unwrap_or_default();
|
||||
let compose_project = labels
|
||||
.get("com.docker.compose.project")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
ContainerInfo {
|
||||
id,
|
||||
name,
|
||||
image: c.image.unwrap_or_default(),
|
||||
status: c.status.unwrap_or_default(),
|
||||
state: c.state.unwrap_or_default(),
|
||||
ports,
|
||||
created_at: c.created.unwrap_or(0),
|
||||
labels,
|
||||
compose_project,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn start(&self, id: &str) -> Result<()> {
|
||||
self.inner
|
||||
.start_container(id, None::<StartContainerOptions<String>>)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop(&self, id: &str) -> Result<()> {
|
||||
self.inner
|
||||
.stop_container(id, Some(StopContainerOptions { t: 10 }))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn restart(&self, id: &str) -> Result<()> {
|
||||
self.inner.restart_container(id, None).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn remove(&self, id: &str) -> Result<()> {
|
||||
self.inner
|
||||
.remove_container(
|
||||
id,
|
||||
Some(RemoveContainerOptions {
|
||||
force: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn logs(
|
||||
&self,
|
||||
id: &str,
|
||||
follow: bool,
|
||||
tail: i32,
|
||||
) -> Pin<Box<dyn Stream<Item = Result<LogOutput, bollard::errors::Error>> + Send>> {
|
||||
let tail_str = if tail > 0 {
|
||||
tail.to_string()
|
||||
} else {
|
||||
"100".to_string()
|
||||
};
|
||||
Box::pin(self.inner.logs(
|
||||
id,
|
||||
Some(LogsOptions::<String> {
|
||||
stdout: true,
|
||||
stderr: true,
|
||||
follow,
|
||||
tail: tail_str,
|
||||
timestamps: false,
|
||||
..Default::default()
|
||||
}),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod tests {
|
||||
use super::*;
|
||||
use bytes::Bytes;
|
||||
use futures_util::StreamExt;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio_stream::once as stream_once;
|
||||
|
||||
// ── Minimal mock backend ──────────────────────────────────────────────────
|
||||
|
||||
/// Records which method was last called and with what container id, so
|
||||
/// tests can assert on behaviour without a real Docker daemon.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct MockBackend {
|
||||
pub calls: Arc<Mutex<Vec<String>>>,
|
||||
/// When Some(msg) every async method returns Err(anyhow!(msg)).
|
||||
pub fail_with: Option<String>,
|
||||
}
|
||||
|
||||
impl MockBackend {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn failing(msg: &str) -> Self {
|
||||
Self {
|
||||
calls: Default::default(),
|
||||
fail_with: Some(msg.to_owned()),
|
||||
}
|
||||
}
|
||||
|
||||
fn record(&self, entry: String) {
|
||||
self.calls.lock().unwrap().push(entry);
|
||||
}
|
||||
|
||||
fn maybe_err(&self) -> Result<()> {
|
||||
if let Some(ref m) = self.fail_with {
|
||||
anyhow::bail!("{}", m);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl ContainerBackend for MockBackend {
|
||||
async fn list_containers(&self) -> Result<Vec<ContainerInfo>> {
|
||||
self.record("list".to_string());
|
||||
self.maybe_err()?;
|
||||
Ok(vec![ContainerInfo {
|
||||
id: "abc123".to_string(),
|
||||
name: "test-container".to_string(),
|
||||
image: "nginx:latest".to_string(),
|
||||
status: "Up 2 hours".to_string(),
|
||||
state: "running".to_string(),
|
||||
ports: vec![ContainerPort {
|
||||
host_port: 8080,
|
||||
container_port: 80,
|
||||
protocol: "tcp".to_string(),
|
||||
host_ip: "0.0.0.0".to_string(),
|
||||
}],
|
||||
created_at: 1_700_000_000,
|
||||
labels: HashMap::new(),
|
||||
compose_project: String::new(),
|
||||
}])
|
||||
}
|
||||
|
||||
async fn start(&self, id: &str) -> Result<()> {
|
||||
self.record(format!("start:{id}"));
|
||||
self.maybe_err()
|
||||
}
|
||||
|
||||
async fn stop(&self, id: &str) -> Result<()> {
|
||||
self.record(format!("stop:{id}"));
|
||||
self.maybe_err()
|
||||
}
|
||||
|
||||
async fn restart(&self, id: &str) -> Result<()> {
|
||||
self.record(format!("restart:{id}"));
|
||||
self.maybe_err()
|
||||
}
|
||||
|
||||
async fn remove(&self, id: &str) -> Result<()> {
|
||||
self.record(format!("remove:{id}"));
|
||||
self.maybe_err()
|
||||
}
|
||||
|
||||
fn logs(
|
||||
&self,
|
||||
id: &str,
|
||||
_follow: bool,
|
||||
_tail: i32,
|
||||
) -> Pin<Box<dyn Stream<Item = Result<LogOutput, bollard::errors::Error>> + Send>>
|
||||
{
|
||||
self.record(format!("logs:{id}"));
|
||||
let chunk = LogOutput::StdOut {
|
||||
message: Bytes::from("hello from mock\n"),
|
||||
};
|
||||
Box::pin(stream_once(Ok(chunk)))
|
||||
}
|
||||
}
|
||||
|
||||
// ── list_containers ───────────────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_list_containers_returns_one_entry() {
|
||||
let backend = MockBackend::new();
|
||||
let containers = backend.list_containers().await.unwrap();
|
||||
assert_eq!(containers.len(), 1);
|
||||
assert_eq!(containers[0].id, "abc123");
|
||||
assert_eq!(containers[0].name, "test-container");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_list_containers_records_call() {
|
||||
let backend = MockBackend::new();
|
||||
backend.list_containers().await.unwrap();
|
||||
let calls = backend.calls.lock().unwrap().clone();
|
||||
assert_eq!(calls, vec!["list"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_list_containers_propagates_error() {
|
||||
let backend = MockBackend::failing("docker down");
|
||||
let err = backend.list_containers().await.unwrap_err();
|
||||
assert!(err.to_string().contains("docker down"));
|
||||
}
|
||||
|
||||
// ── start / stop / restart / remove ──────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_start_records_id() {
|
||||
let backend = MockBackend::new();
|
||||
backend.start("cid-1").await.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["start:cid-1"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_stop_records_id() {
|
||||
let backend = MockBackend::new();
|
||||
backend.stop("cid-2").await.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["stop:cid-2"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_restart_records_id() {
|
||||
let backend = MockBackend::new();
|
||||
backend.restart("cid-3").await.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["restart:cid-3"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_remove_records_id() {
|
||||
let backend = MockBackend::new();
|
||||
backend.remove("cid-4").await.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["remove:cid-4"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_operations_propagate_errors() {
|
||||
let backend = MockBackend::failing("socket gone");
|
||||
assert!(backend.start("x").await.is_err());
|
||||
assert!(backend.stop("x").await.is_err());
|
||||
assert!(backend.restart("x").await.is_err());
|
||||
assert!(backend.remove("x").await.is_err());
|
||||
}
|
||||
|
||||
// ── logs stream ──────────────────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_logs_yields_stdout_chunk() {
|
||||
let backend = MockBackend::new();
|
||||
let mut stream = backend.logs("cid-5", false, 10);
|
||||
let item = stream.next().await.unwrap().unwrap();
|
||||
match item {
|
||||
LogOutput::StdOut { message } => {
|
||||
assert_eq!(message.as_ref(), b"hello from mock\n");
|
||||
}
|
||||
other => panic!("unexpected variant: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mock_logs_records_id() {
|
||||
let backend = MockBackend::new();
|
||||
let mut stream = backend.logs("cid-5", false, 10);
|
||||
// drain the stream
|
||||
while stream.next().await.is_some() {}
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["logs:cid-5"]);
|
||||
}
|
||||
|
||||
// ── ContainerInfo / ContainerPort field mapping ───────────────────────────
|
||||
|
||||
#[test]
|
||||
fn container_info_fields_are_accessible() {
|
||||
let port = ContainerPort {
|
||||
host_port: 443,
|
||||
container_port: 8443,
|
||||
protocol: "tcp".to_string(),
|
||||
host_ip: "127.0.0.1".to_string(),
|
||||
};
|
||||
assert_eq!(port.host_port, 443);
|
||||
assert_eq!(port.container_port, 8443);
|
||||
assert_eq!(port.protocol, "tcp");
|
||||
|
||||
let info = ContainerInfo {
|
||||
id: "id1".to_string(),
|
||||
name: "name1".to_string(),
|
||||
image: "img".to_string(),
|
||||
status: "running".to_string(),
|
||||
state: "running".to_string(),
|
||||
ports: vec![port],
|
||||
created_at: 42,
|
||||
labels: HashMap::new(),
|
||||
compose_project: "proj".to_string(),
|
||||
};
|
||||
assert_eq!(info.ports.len(), 1);
|
||||
assert_eq!(info.compose_project, "proj");
|
||||
}
|
||||
}
|
||||
434
agent/src/main.rs
Normal file
434
agent/src/main.rs
Normal file
@ -0,0 +1,434 @@
|
||||
mod docker;
|
||||
|
||||
pub mod proto {
|
||||
tonic::include_proto!("containarr.agent.v1");
|
||||
}
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use bollard::container::LogOutput;
|
||||
use docker::{ContainerBackend, DockerClient};
|
||||
use futures_util::StreamExt as _;
|
||||
use proto::{
|
||||
agent_gateway_client::AgentGatewayClient,
|
||||
agent_message, server_message,
|
||||
AgentHandshake, AgentMessage, ContainerAction, ContainerSnapshot,
|
||||
};
|
||||
use std::{collections::HashMap, env, time::Duration};
|
||||
use tokio::{sync::mpsc, task::JoinHandle, time};
|
||||
use tonic::Request;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
const SNAPSHOT_INTERVAL: Duration = Duration::from_secs(10);
|
||||
const RECONNECT_DELAY: Duration = Duration::from_secs(5);
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.json()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "info".into()),
|
||||
)
|
||||
.init();
|
||||
|
||||
let server_url = env::var("CONTAINARR_SERVER_URL")
|
||||
.context("CONTAINARR_SERVER_URL not set")?;
|
||||
let token = env::var("CONTAINARR_AGENT_TOKEN")
|
||||
.context("CONTAINARR_AGENT_TOKEN not set")?;
|
||||
let hostname = env::var("HOSTNAME").unwrap_or_else(|_| "unknown".into());
|
||||
|
||||
let docker = DockerClient::new().context("connect to Docker socket")?;
|
||||
|
||||
loop {
|
||||
if let Err(e) = run(&server_url, &token, &hostname, docker.clone()).await {
|
||||
error!("connection lost: {:#}", e);
|
||||
}
|
||||
info!("reconnecting in {:?}", RECONNECT_DELAY);
|
||||
time::sleep(RECONNECT_DELAY).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn run(url: &str, token: &str, hostname: &str, docker: DockerClient) -> Result<()> {
|
||||
info!("connecting to server at {}", url);
|
||||
let mut client = AgentGatewayClient::connect(url.to_string()).await?;
|
||||
|
||||
let (tx, rx) = mpsc::channel::<AgentMessage>(64);
|
||||
|
||||
tx.send(AgentMessage {
|
||||
payload: Some(agent_message::Payload::Handshake(AgentHandshake {
|
||||
token: token.to_string(),
|
||||
hostname: hostname.to_string(),
|
||||
arch: std::env::consts::ARCH.to_string(),
|
||||
os: std::env::consts::OS.to_string(),
|
||||
})),
|
||||
})
|
||||
.await?;
|
||||
|
||||
let outbound = tokio_stream::wrappers::ReceiverStream::new(rx);
|
||||
let mut inbound = client.tunnel(Request::new(outbound)).await?.into_inner();
|
||||
|
||||
let mut snapshot_ticker = time::interval(SNAPSHOT_INTERVAL);
|
||||
let mut log_tasks: HashMap<String, JoinHandle<()>> = HashMap::new();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = snapshot_ticker.tick() => {
|
||||
match time::timeout(Duration::from_secs(5), docker.list_containers()).await {
|
||||
Err(_) => warn!("docker list timed out"),
|
||||
Ok(Err(e)) => warn!("docker list failed: {:#}", e),
|
||||
Ok(Ok(containers)) => {
|
||||
let msg = AgentMessage {
|
||||
payload: Some(agent_message::Payload::Snapshot(ContainerSnapshot {
|
||||
containers,
|
||||
timestamp: unix_now(),
|
||||
})),
|
||||
};
|
||||
if tx.send(msg).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result = inbound.message() => {
|
||||
match result? {
|
||||
None => break,
|
||||
Some(msg) => match msg.payload {
|
||||
Some(server_message::Payload::ContainerCmd(cmd)) => {
|
||||
let ok = execute_action(&docker, &cmd.container_id, cmd.action).await;
|
||||
let _ = tx.send(AgentMessage {
|
||||
payload: Some(agent_message::Payload::Result(proto::CommandResult {
|
||||
command_id: cmd.command_id,
|
||||
success: ok.is_ok(),
|
||||
error: ok.err().map(|e| e.to_string()).unwrap_or_default(),
|
||||
})),
|
||||
}).await;
|
||||
}
|
||||
Some(server_message::Payload::StreamLogs(cmd)) => {
|
||||
if let Some(old) = log_tasks.remove(&cmd.container_id) {
|
||||
old.abort();
|
||||
}
|
||||
let docker_clone = docker.clone();
|
||||
let tx_clone = tx.clone();
|
||||
let cid = cmd.container_id.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
let mut stream = docker_clone.logs(&cid, cmd.follow, cmd.tail);
|
||||
while let Some(result) = stream.next().await {
|
||||
match result {
|
||||
Ok(output) => {
|
||||
let (stream_name, data) = match output {
|
||||
LogOutput::StdOut { message } => ("stdout", message),
|
||||
LogOutput::StdErr { message } => ("stderr", message),
|
||||
_ => continue,
|
||||
};
|
||||
let msg = AgentMessage {
|
||||
payload: Some(agent_message::Payload::LogChunk(
|
||||
proto::LogChunk {
|
||||
container_id: cid.clone(),
|
||||
stream: stream_name.to_string(),
|
||||
data: data.to_vec(),
|
||||
timestamp: unix_now(),
|
||||
},
|
||||
)),
|
||||
};
|
||||
if tx_clone.send(msg).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("log stream error for {}: {:#}", cid, e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
log_tasks.insert(cmd.container_id, handle);
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (_, task) in log_tasks {
|
||||
task.abort();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_action<B: ContainerBackend>(
|
||||
docker: &B,
|
||||
id: &str,
|
||||
action: i32,
|
||||
) -> Result<()> {
|
||||
match ContainerAction::try_from(action)? {
|
||||
ContainerAction::Start => docker.start(id).await,
|
||||
ContainerAction::Stop => docker.stop(id).await,
|
||||
ContainerAction::Restart => docker.restart(id).await,
|
||||
ContainerAction::Remove => docker.remove(id).await,
|
||||
ContainerAction::Unspecified => anyhow::bail!("unspecified action"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn unix_now() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs() as i64
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::docker::tests::MockBackend;
|
||||
use proto::ContainerAction;
|
||||
|
||||
// ── unix_now ──────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn unix_now_is_positive_and_recent() {
|
||||
let t = unix_now();
|
||||
// Should be somewhere past 2020-01-01 (1_577_836_800)
|
||||
assert!(t > 1_577_836_800, "unix_now returned {t}, expected > 2020");
|
||||
// And not absurdly far in the future (year 2100 = 4_102_444_800)
|
||||
assert!(t < 4_102_444_800, "unix_now returned {t}, looks wrong");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unix_now_is_monotone_or_equal() {
|
||||
let t1 = unix_now();
|
||||
let t2 = unix_now();
|
||||
// Same second or the second just ticked forward — never backwards
|
||||
assert!(t2 >= t1);
|
||||
}
|
||||
|
||||
// ── ContainerAction enum parsing ──────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn container_action_from_valid_int() {
|
||||
assert_eq!(ContainerAction::try_from(0).unwrap(), ContainerAction::Unspecified);
|
||||
assert_eq!(ContainerAction::try_from(1).unwrap(), ContainerAction::Start);
|
||||
assert_eq!(ContainerAction::try_from(2).unwrap(), ContainerAction::Stop);
|
||||
assert_eq!(ContainerAction::try_from(3).unwrap(), ContainerAction::Restart);
|
||||
assert_eq!(ContainerAction::try_from(4).unwrap(), ContainerAction::Remove);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn container_action_from_invalid_int_errors() {
|
||||
assert!(ContainerAction::try_from(99).is_err());
|
||||
assert!(ContainerAction::try_from(-1).is_err());
|
||||
}
|
||||
|
||||
// ── execute_action — routing ──────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_start_calls_start() {
|
||||
let backend = MockBackend::new();
|
||||
execute_action(&backend, "container-a", ContainerAction::Start as i32)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["start:container-a"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_stop_calls_stop() {
|
||||
let backend = MockBackend::new();
|
||||
execute_action(&backend, "container-b", ContainerAction::Stop as i32)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["stop:container-b"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_restart_calls_restart() {
|
||||
let backend = MockBackend::new();
|
||||
execute_action(&backend, "container-c", ContainerAction::Restart as i32)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["restart:container-c"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_remove_calls_remove() {
|
||||
let backend = MockBackend::new();
|
||||
execute_action(&backend, "container-d", ContainerAction::Remove as i32)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(*backend.calls.lock().unwrap(), vec!["remove:container-d"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_unspecified_returns_error() {
|
||||
let backend = MockBackend::new();
|
||||
let result =
|
||||
execute_action(&backend, "container-e", ContainerAction::Unspecified as i32).await;
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("unspecified action"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_invalid_int_returns_error() {
|
||||
let backend = MockBackend::new();
|
||||
let result = execute_action(&backend, "container-f", 999).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
// ── execute_action — error propagation ───────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_start_propagates_backend_error() {
|
||||
let backend = MockBackend::failing("docker not reachable");
|
||||
let result =
|
||||
execute_action(&backend, "x", ContainerAction::Start as i32).await;
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("docker not reachable"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_stop_propagates_backend_error() {
|
||||
let backend = MockBackend::failing("timeout");
|
||||
let result =
|
||||
execute_action(&backend, "x", ContainerAction::Stop as i32).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_restart_propagates_backend_error() {
|
||||
let backend = MockBackend::failing("timeout");
|
||||
let result =
|
||||
execute_action(&backend, "x", ContainerAction::Restart as i32).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_action_remove_propagates_backend_error() {
|
||||
let backend = MockBackend::failing("permission denied");
|
||||
let result =
|
||||
execute_action(&backend, "x", ContainerAction::Remove as i32).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
// ── Log-task abort logic (unit-level) ─────────────────────────────────────
|
||||
|
||||
/// Spawn a long-running task and verify it is aborted when its JoinHandle
|
||||
/// is dropped via abort() — exercises the same flow as the StreamLogs branch.
|
||||
#[tokio::test]
|
||||
async fn spawned_task_can_be_aborted() {
|
||||
use tokio::sync::oneshot;
|
||||
let (started_tx, started_rx) = oneshot::channel::<()>();
|
||||
let handle = tokio::spawn(async move {
|
||||
let _ = started_tx.send(());
|
||||
// Block indefinitely simulating a live log stream
|
||||
tokio::time::sleep(Duration::from_secs(3600)).await;
|
||||
});
|
||||
// Wait until the task is definitely running
|
||||
started_rx.await.unwrap();
|
||||
handle.abort();
|
||||
let result = handle.await;
|
||||
assert!(result.is_err()); // JoinError with is_cancelled() == true
|
||||
assert!(result.unwrap_err().is_cancelled());
|
||||
}
|
||||
|
||||
/// Spawning a second task for the same container_id aborts the first one,
|
||||
/// mirroring the log_tasks.remove() + old.abort() pattern in main.rs.
|
||||
#[tokio::test]
|
||||
async fn second_stream_aborts_first() {
|
||||
use tokio::sync::oneshot;
|
||||
let mut log_tasks: HashMap<String, JoinHandle<()>> = HashMap::new();
|
||||
let cid = "my-container".to_string();
|
||||
|
||||
let (tx1, rx1) = oneshot::channel::<()>();
|
||||
let h1 = tokio::spawn(async move {
|
||||
let _ = tx1.send(());
|
||||
tokio::time::sleep(Duration::from_secs(3600)).await;
|
||||
});
|
||||
rx1.await.unwrap();
|
||||
log_tasks.insert(cid.clone(), h1);
|
||||
|
||||
// Second StreamLogs for the same cid — abort the old task
|
||||
if let Some(old) = log_tasks.remove(&cid) {
|
||||
old.abort();
|
||||
}
|
||||
let (tx2, rx2) = oneshot::channel::<()>();
|
||||
let h2 = tokio::spawn(async move {
|
||||
let _ = tx2.send(());
|
||||
tokio::time::sleep(Duration::from_secs(3600)).await;
|
||||
});
|
||||
rx2.await.unwrap();
|
||||
log_tasks.insert(cid.clone(), h2);
|
||||
|
||||
assert_eq!(log_tasks.len(), 1);
|
||||
// Clean up
|
||||
for (_, h) in log_tasks {
|
||||
h.abort();
|
||||
}
|
||||
}
|
||||
|
||||
// ── Proto message construction ────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn agent_handshake_fields_roundtrip() {
|
||||
let hs = AgentHandshake {
|
||||
token: "tok".to_string(),
|
||||
hostname: "host".to_string(),
|
||||
arch: "x86_64".to_string(),
|
||||
os: "linux".to_string(),
|
||||
};
|
||||
assert_eq!(hs.token, "tok");
|
||||
assert_eq!(hs.hostname, "host");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_message_wraps_handshake() {
|
||||
let msg = AgentMessage {
|
||||
payload: Some(agent_message::Payload::Handshake(AgentHandshake {
|
||||
token: "t".to_string(),
|
||||
hostname: "h".to_string(),
|
||||
arch: "arm64".to_string(),
|
||||
os: "linux".to_string(),
|
||||
})),
|
||||
};
|
||||
assert!(matches!(
|
||||
msg.payload,
|
||||
Some(agent_message::Payload::Handshake(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn command_result_ok_fields() {
|
||||
let r = proto::CommandResult {
|
||||
command_id: "cmd-1".to_string(),
|
||||
success: true,
|
||||
error: String::new(),
|
||||
};
|
||||
assert!(r.success);
|
||||
assert!(r.error.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn command_result_err_fields() {
|
||||
let r = proto::CommandResult {
|
||||
command_id: "cmd-2".to_string(),
|
||||
success: false,
|
||||
error: "container not found".to_string(),
|
||||
};
|
||||
assert!(!r.success);
|
||||
assert_eq!(r.error, "container not found");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn log_chunk_fields() {
|
||||
let chunk = proto::LogChunk {
|
||||
container_id: "cid".to_string(),
|
||||
stream: "stdout".to_string(),
|
||||
data: b"hello".to_vec(),
|
||||
timestamp: 12345,
|
||||
};
|
||||
assert_eq!(chunk.stream, "stdout");
|
||||
assert_eq!(chunk.data, b"hello");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user