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::{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
75#[async_trait]
76impl Runtime for MockRuntime {
77 fn name(&self) -> &str {
78 "mock"
79 }
80
81 async fn create(&self, spec: &WorkloadSpec) -> Result<WorkloadHandle> {
82 let mut counter = self.counter.lock().await;
83 *counter += 1;
84 let id = format!("mock-{}", *counter);
85
86 self.ops
87 .lock()
88 .await
89 .push(MockOp::Create(spec.name.clone()));
90 self.statuses
91 .lock()
92 .await
93 .insert(id.clone(), WorkloadStatus::Creating);
94
95 Ok(WorkloadHandle {
96 runtime_id: id,
97 name: format!("orca-{}", spec.name),
98 metadata: HashMap::new(),
99 })
100 }
101
102 async fn start(&self, handle: &WorkloadHandle) -> Result<()> {
103 self.ops
104 .lock()
105 .await
106 .push(MockOp::Start(handle.name.clone()));
107 self.statuses
108 .lock()
109 .await
110 .insert(handle.runtime_id.clone(), WorkloadStatus::Running);
111 Ok(())
112 }
113
114 async fn stop(&self, handle: &WorkloadHandle, _timeout: Duration) -> Result<()> {
115 self.ops
116 .lock()
117 .await
118 .push(MockOp::Stop(handle.name.clone()));
119 self.statuses
120 .lock()
121 .await
122 .insert(handle.runtime_id.clone(), WorkloadStatus::Stopped);
123 Ok(())
124 }
125
126 async fn remove(&self, handle: &WorkloadHandle) -> Result<()> {
127 self.ops
128 .lock()
129 .await
130 .push(MockOp::Remove(handle.name.clone()));
131 self.statuses.lock().await.remove(&handle.runtime_id);
132 Ok(())
133 }
134
135 async fn status(&self, handle: &WorkloadHandle) -> Result<WorkloadStatus> {
136 let statuses = self.statuses.lock().await;
137 statuses
138 .get(&handle.runtime_id)
139 .copied()
140 .ok_or_else(|| OrcaError::WorkloadNotFound {
141 name: handle.runtime_id.clone(),
142 })
143 }
144
145 async fn logs(&self, _handle: &WorkloadHandle, _opts: &LogOpts) -> Result<LogStream> {
146 let text = b"mock log line 1\nmock log line 2\n";
147 let cursor = std::io::Cursor::new(text.to_vec());
148 Ok(Box::pin(cursor) as Pin<Box<dyn tokio::io::AsyncRead + Send>>)
149 }
150
151 async fn exec(&self, _handle: &WorkloadHandle, cmd: &[String]) -> Result<ExecResult> {
152 Ok(ExecResult {
153 exit_code: 0,
154 stdout: format!("mock exec: {}", cmd.join(" ")).into_bytes(),
155 stderr: Vec::new(),
156 })
157 }
158
159 async fn stats(&self, _handle: &WorkloadHandle) -> Result<ResourceStats> {
160 Ok(ResourceStats {
161 cpu_percent: 0.0,
162 memory_bytes: 0,
163 network_rx_bytes: 0,
164 network_tx_bytes: 0,
165 gpu_stats: Vec::new(),
166 timestamp: Utc::now(),
167 })
168 }
169
170 async fn resolve_host_port(
171 &self,
172 _handle: &WorkloadHandle,
173 _container_port: u16,
174 ) -> Result<Option<u16>> {
175 Ok(self.mock_host_port)
176 }
177}