use std::collections::HashMap;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use futures::stream::{Stream, StreamExt};
use crate::error::RuntimeError;
use crate::runtime::{ContainerId, ContainerRuntime, ContainerStatus, LogChunk, LogChunkStream};
use lightshuttle_spec::ContainerSpec;
#[derive(Clone)]
pub struct MockRuntime {
state: Arc<Mutex<HashMap<String, MockContainer>>>,
fail_on: Arc<Mutex<Option<String>>>,
start_order: Arc<Mutex<Vec<String>>>,
stop_order: Arc<Mutex<Vec<String>>>,
started_specs: Arc<Mutex<Vec<ContainerSpec>>>,
}
struct MockContainer {
name: String,
status: ContainerStatus,
started_at: Instant,
healthy_after: Duration,
}
impl MockRuntime {
#[must_use]
pub fn new() -> Self {
Self {
state: Arc::new(Mutex::new(HashMap::new())),
fail_on: Arc::new(Mutex::new(None)),
start_order: Arc::new(Mutex::new(Vec::new())),
stop_order: Arc::new(Mutex::new(Vec::new())),
started_specs: Arc::new(Mutex::new(Vec::new())),
}
}
pub fn fail_on(&self, name: &str) {
*self.fail_on.lock().expect("fail_on mutex poisoned") = Some(name.to_owned());
}
#[must_use]
pub fn started_resources(&self) -> Vec<String> {
self.start_order
.lock()
.expect("start_order mutex poisoned")
.clone()
}
#[must_use]
pub fn stopped_resources(&self) -> Vec<String> {
self.stop_order
.lock()
.expect("stop_order mutex poisoned")
.clone()
}
#[must_use]
pub fn started_specs(&self) -> Vec<ContainerSpec> {
self.started_specs
.lock()
.expect("started_specs mutex poisoned")
.clone()
}
}
impl Default for MockRuntime {
fn default() -> Self {
Self::new()
}
}
impl ContainerRuntime for MockRuntime {
async fn start(&self, spec: &ContainerSpec) -> Result<ContainerId, RuntimeError> {
if self
.fail_on
.lock()
.expect("fail_on mutex poisoned")
.as_deref()
== Some(spec.name.as_str())
{
return Err(RuntimeError::InvalidSpec(format!(
"mock failure for `{}`",
spec.name
)));
}
let id = ContainerId::new(format!("mock-{}", spec.name));
if self
.state
.lock()
.expect("state mutex poisoned")
.contains_key(id.as_str())
{
return Err(RuntimeError::InvalidSpec(format!(
"container name `{}` already in use",
spec.name
)));
}
self.start_order
.lock()
.expect("start_order mutex poisoned")
.push(spec.name.clone());
self.started_specs
.lock()
.expect("started_specs mutex poisoned")
.push(spec.clone());
self.state.lock().expect("state mutex poisoned").insert(
id.as_str().to_owned(),
MockContainer {
name: spec.name.clone(),
status: ContainerStatus::Starting,
started_at: Instant::now(),
healthy_after: Duration::from_millis(30),
},
);
Ok(id)
}
async fn stop(&self, id: &ContainerId, _grace: Duration) -> Result<(), RuntimeError> {
let mut state = self.state.lock().expect("state mutex poisoned");
if let Some(c) = state.get_mut(id.as_str()) {
c.status = ContainerStatus::Stopped { exit_code: Some(0) };
self.stop_order
.lock()
.expect("stop_order mutex poisoned")
.push(c.name.clone());
}
Ok(())
}
async fn remove(&self, name: &str) -> Result<(), RuntimeError> {
self.state
.lock()
.expect("state mutex poisoned")
.remove(&format!("mock-{name}"));
Ok(())
}
async fn inspect(&self, id: &ContainerId) -> Result<ContainerStatus, RuntimeError> {
let state = self.state.lock().expect("state mutex poisoned");
let c = state
.get(id.as_str())
.ok_or_else(|| RuntimeError::NotFound(id.as_str().to_owned()))?;
Ok(c.status.clone())
}
async fn wait_healthy(&self, id: &ContainerId, timeout: Duration) -> Result<(), RuntimeError> {
let start = Instant::now();
while start.elapsed() < timeout {
{
let mut state = self.state.lock().expect("state mutex poisoned");
if let Some(c) = state.get_mut(id.as_str())
&& c.started_at.elapsed() >= c.healthy_after
{
c.status = ContainerStatus::Healthy;
return Ok(());
}
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
Err(RuntimeError::Timeout {
operation: "wait_healthy",
after: timeout,
})
}
async fn logs(&self, _id: &ContainerId, _follow: bool) -> Result<LogChunkStream, RuntimeError> {
let empty: Pin<Box<dyn Stream<Item = Result<LogChunk, RuntimeError>> + Send>> =
Box::pin(futures::stream::empty::<Result<LogChunk, RuntimeError>>().map(|x| x));
Ok(empty)
}
async fn ensure_project_network(&self, _project: &str) -> Result<(), RuntimeError> {
Ok(())
}
async fn teardown_project_network(&self, _project: &str) -> Result<(), RuntimeError> {
Ok(())
}
}