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::{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
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}