lightshuttle_runtime/
testkit.rs1use std::collections::HashMap;
42use std::pin::Pin;
43use std::sync::{Arc, Mutex};
44use std::time::{Duration, Instant};
45
46use futures::stream::{Stream, StreamExt};
47
48use crate::error::RuntimeError;
49use crate::runtime::{ContainerId, ContainerRuntime, ContainerStatus, LogChunk, LogChunkStream};
50use lightshuttle_spec::ContainerSpec;
51
52#[derive(Clone)]
58pub struct MockRuntime {
59 state: Arc<Mutex<HashMap<String, MockContainer>>>,
60 fail_on: Arc<Mutex<Option<String>>>,
61 start_order: Arc<Mutex<Vec<String>>>,
62 stop_order: Arc<Mutex<Vec<String>>>,
63 started_specs: Arc<Mutex<Vec<ContainerSpec>>>,
64}
65
66struct MockContainer {
67 name: String,
68 status: ContainerStatus,
69 started_at: Instant,
70 healthy_after: Duration,
71}
72
73impl MockRuntime {
74 #[must_use]
76 pub fn new() -> Self {
77 Self {
78 state: Arc::new(Mutex::new(HashMap::new())),
79 fail_on: Arc::new(Mutex::new(None)),
80 start_order: Arc::new(Mutex::new(Vec::new())),
81 stop_order: Arc::new(Mutex::new(Vec::new())),
82 started_specs: Arc::new(Mutex::new(Vec::new())),
83 }
84 }
85
86 pub fn fail_on(&self, name: &str) {
92 *self.fail_on.lock().expect("fail_on mutex poisoned") = Some(name.to_owned());
93 }
94
95 #[must_use]
97 pub fn started_resources(&self) -> Vec<String> {
98 self.start_order
99 .lock()
100 .expect("start_order mutex poisoned")
101 .clone()
102 }
103
104 #[must_use]
106 pub fn stopped_resources(&self) -> Vec<String> {
107 self.stop_order
108 .lock()
109 .expect("stop_order mutex poisoned")
110 .clone()
111 }
112
113 #[must_use]
115 pub fn started_specs(&self) -> Vec<ContainerSpec> {
116 self.started_specs
117 .lock()
118 .expect("started_specs mutex poisoned")
119 .clone()
120 }
121}
122
123impl Default for MockRuntime {
124 fn default() -> Self {
125 Self::new()
126 }
127}
128
129impl ContainerRuntime for MockRuntime {
130 async fn start(&self, spec: &ContainerSpec) -> Result<ContainerId, RuntimeError> {
131 if self
132 .fail_on
133 .lock()
134 .expect("fail_on mutex poisoned")
135 .as_deref()
136 == Some(spec.name.as_str())
137 {
138 return Err(RuntimeError::InvalidSpec(format!(
139 "mock failure for `{}`",
140 spec.name
141 )));
142 }
143 let id = ContainerId::new(format!("mock-{}", spec.name));
144 if self
145 .state
146 .lock()
147 .expect("state mutex poisoned")
148 .contains_key(id.as_str())
149 {
150 return Err(RuntimeError::InvalidSpec(format!(
151 "container name `{}` already in use",
152 spec.name
153 )));
154 }
155 self.start_order
156 .lock()
157 .expect("start_order mutex poisoned")
158 .push(spec.name.clone());
159 self.started_specs
160 .lock()
161 .expect("started_specs mutex poisoned")
162 .push(spec.clone());
163 self.state.lock().expect("state mutex poisoned").insert(
164 id.as_str().to_owned(),
165 MockContainer {
166 name: spec.name.clone(),
167 status: ContainerStatus::Starting,
168 started_at: Instant::now(),
169 healthy_after: Duration::from_millis(30),
170 },
171 );
172 Ok(id)
173 }
174
175 async fn stop(&self, id: &ContainerId, _grace: Duration) -> Result<(), RuntimeError> {
176 let mut state = self.state.lock().expect("state mutex poisoned");
177 if let Some(c) = state.get_mut(id.as_str()) {
178 c.status = ContainerStatus::Stopped { exit_code: Some(0) };
179 self.stop_order
180 .lock()
181 .expect("stop_order mutex poisoned")
182 .push(c.name.clone());
183 }
184 Ok(())
185 }
186
187 async fn remove(&self, name: &str) -> Result<(), RuntimeError> {
188 self.state
189 .lock()
190 .expect("state mutex poisoned")
191 .remove(&format!("mock-{name}"));
192 Ok(())
193 }
194
195 async fn inspect(&self, id: &ContainerId) -> Result<ContainerStatus, RuntimeError> {
196 let state = self.state.lock().expect("state mutex poisoned");
197 let c = state
198 .get(id.as_str())
199 .ok_or_else(|| RuntimeError::NotFound(id.as_str().to_owned()))?;
200 Ok(c.status.clone())
201 }
202
203 async fn wait_healthy(&self, id: &ContainerId, timeout: Duration) -> Result<(), RuntimeError> {
204 let start = Instant::now();
205 while start.elapsed() < timeout {
206 {
207 let mut state = self.state.lock().expect("state mutex poisoned");
208 if let Some(c) = state.get_mut(id.as_str())
209 && c.started_at.elapsed() >= c.healthy_after
210 {
211 c.status = ContainerStatus::Healthy;
212 return Ok(());
213 }
214 }
215 tokio::time::sleep(Duration::from_millis(10)).await;
216 }
217 Err(RuntimeError::Timeout {
218 operation: "wait_healthy",
219 after: timeout,
220 })
221 }
222
223 async fn logs(&self, _id: &ContainerId, _follow: bool) -> Result<LogChunkStream, RuntimeError> {
224 let empty: Pin<Box<dyn Stream<Item = Result<LogChunk, RuntimeError>> + Send>> =
225 Box::pin(futures::stream::empty::<Result<LogChunk, RuntimeError>>().map(|x| x));
226 Ok(empty)
227 }
228
229 async fn ensure_project_network(&self, _project: &str) -> Result<(), RuntimeError> {
230 Ok(())
231 }
232
233 async fn teardown_project_network(&self, _project: &str) -> Result<(), RuntimeError> {
234 Ok(())
235 }
236}