orca_core/testing/
mock_runtime.rs1use std::collections::HashMap;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::time::Duration;
7
8use async_trait::async_trait;
9use chrono::Utc;
10use tokio::sync::Mutex;
11
12use crate::error::{OrcaError, Result};
13use crate::runtime::{AsAny, ExecResult, LogOpts, LogStream, Runtime, WorkloadHandle};
14use crate::types::{ResourceStats, WorkloadSpec, WorkloadStatus};
15
16#[derive(Debug, Clone)]
18pub enum MockOp {
19 Create(String),
21 Start(String),
23 Stop(String),
25 Remove(String),
27}
28
29pub struct MockRuntime {
34 pub ops: Arc<Mutex<Vec<MockOp>>>,
36 statuses: Arc<Mutex<HashMap<String, WorkloadStatus>>>,
38 counter: Arc<Mutex<u64>>,
40 pub mock_host_port: Option<u16>,
42}
43
44impl MockRuntime {
45 pub fn new() -> Self {
47 Self {
48 ops: Arc::new(Mutex::new(Vec::new())),
49 statuses: Arc::new(Mutex::new(HashMap::new())),
50 counter: Arc::new(Mutex::new(0)),
51 mock_host_port: None,
52 }
53 }
54
55 pub fn with_host_port(port: u16) -> Self {
57 Self {
58 mock_host_port: Some(port),
59 ..Self::new()
60 }
61 }
62
63 pub async fn recorded_ops(&self) -> Vec<MockOp> {
65 self.ops.lock().await.clone()
66 }
67}
68
69impl Default for MockRuntime {
70 fn default() -> Self {
71 Self::new()
72 }
73}
74
75impl AsAny for MockRuntime {
76 fn as_any(&self) -> &dyn std::any::Any {
77 self
78 }
79}
80
81#[async_trait]
82impl Runtime for MockRuntime {
83 fn name(&self) -> &str {
84 "mock"
85 }
86
87 async fn create(&self, spec: &WorkloadSpec) -> Result<WorkloadHandle> {
88 let mut counter = self.counter.lock().await;
89 *counter += 1;
90 let id = format!("mock-{}", *counter);
91
92 self.ops
93 .lock()
94 .await
95 .push(MockOp::Create(spec.name.clone()));
96 self.statuses
97 .lock()
98 .await
99 .insert(id.clone(), WorkloadStatus::Creating);
100
101 Ok(WorkloadHandle {
102 runtime_id: id,
103 name: format!("orca-{}", spec.name),
104 metadata: HashMap::new(),
105 })
106 }
107
108 async fn start(&self, handle: &WorkloadHandle) -> Result<()> {
109 self.ops
110 .lock()
111 .await
112 .push(MockOp::Start(handle.name.clone()));
113 self.statuses
114 .lock()
115 .await
116 .insert(handle.runtime_id.clone(), WorkloadStatus::Running);
117 Ok(())
118 }
119
120 async fn stop(&self, handle: &WorkloadHandle, _timeout: Duration) -> Result<()> {
121 self.ops
122 .lock()
123 .await
124 .push(MockOp::Stop(handle.name.clone()));
125 self.statuses
126 .lock()
127 .await
128 .insert(handle.runtime_id.clone(), WorkloadStatus::Stopped);
129 Ok(())
130 }
131
132 async fn remove(&self, handle: &WorkloadHandle) -> Result<()> {
133 self.ops
134 .lock()
135 .await
136 .push(MockOp::Remove(handle.name.clone()));
137 self.statuses.lock().await.remove(&handle.runtime_id);
138 Ok(())
139 }
140
141 async fn status(&self, handle: &WorkloadHandle) -> Result<WorkloadStatus> {
142 let statuses = self.statuses.lock().await;
143 statuses
144 .get(&handle.runtime_id)
145 .copied()
146 .ok_or_else(|| OrcaError::WorkloadNotFound {
147 name: handle.runtime_id.clone(),
148 })
149 }
150
151 async fn logs(&self, _handle: &WorkloadHandle, _opts: &LogOpts) -> Result<LogStream> {
152 let text = b"mock log line 1\nmock log line 2\n";
153 let cursor = std::io::Cursor::new(text.to_vec());
154 Ok(Box::pin(cursor) as Pin<Box<dyn tokio::io::AsyncRead + Send>>)
155 }
156
157 async fn exec(&self, _handle: &WorkloadHandle, cmd: &[String]) -> Result<ExecResult> {
158 Ok(ExecResult {
159 exit_code: 0,
160 stdout: format!("mock exec: {}", cmd.join(" ")).into_bytes(),
161 stderr: Vec::new(),
162 })
163 }
164
165 async fn stats(&self, _handle: &WorkloadHandle) -> Result<ResourceStats> {
166 Ok(ResourceStats {
167 cpu_percent: 0.0,
168 memory_bytes: 0,
169 network_rx_bytes: 0,
170 network_tx_bytes: 0,
171 gpu_stats: Vec::new(),
172 timestamp: Utc::now(),
173 })
174 }
175
176 async fn resolve_host_port(
177 &self,
178 _handle: &WorkloadHandle,
179 _container_port: u16,
180 ) -> Result<Option<u16>> {
181 Ok(self.mock_host_port)
182 }
183}