Skip to main content

orca_core/testing/
mock_runtime.rs

1//! Mock implementation of the [`Runtime`] trait for testing.
2
3use 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/// Records of operations performed on the mock runtime.
17#[derive(Debug, Clone)]
18pub enum MockOp {
19    /// A workload was created.
20    Create(String),
21    /// A workload was started.
22    Start(String),
23    /// A workload was stopped.
24    Stop(String),
25    /// A workload was removed.
26    Remove(String),
27}
28
29/// A mock [`Runtime`] that tracks operations without running real workloads.
30///
31/// Use this in integration tests to verify reconciler behavior,
32/// API endpoints, and other components that depend on a runtime.
33pub struct MockRuntime {
34    /// Recorded operations, in order.
35    pub ops: Arc<Mutex<Vec<MockOp>>>,
36    /// Current status per runtime_id.
37    statuses: Arc<Mutex<HashMap<String, WorkloadStatus>>>,
38    /// Counter for generating unique IDs.
39    counter: Arc<Mutex<u64>>,
40    /// If set, the mock host port returned by resolve_host_port.
41    pub mock_host_port: Option<u16>,
42}
43
44impl MockRuntime {
45    /// Create a new mock runtime.
46    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    /// Create a mock runtime that returns a fixed host port.
56    pub fn with_host_port(port: u16) -> Self {
57        Self {
58            mock_host_port: Some(port),
59            ..Self::new()
60        }
61    }
62
63    /// Get a copy of all recorded operations.
64    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}