forge_core/testing/context/
daemon.rs

1//! Test context for daemon functions.
2
3#![allow(clippy::unwrap_used, clippy::indexing_slicing)]
4
5use std::collections::HashMap;
6use std::sync::Arc;
7
8use sqlx::PgPool;
9use tokio::sync::{Mutex, watch};
10use uuid::Uuid;
11
12use super::super::mock_http::{MockHttp, MockRequest, MockResponse};
13use crate::Result;
14use crate::env::{EnvAccess, EnvProvider, MockEnvProvider};
15
16/// Test context for daemon functions.
17pub struct TestDaemonContext {
18    pub daemon_name: String,
19    pub instance_id: Uuid,
20    pool: Option<PgPool>,
21    http: Arc<MockHttp>,
22    /// For triggering shutdown from the test side.
23    pub shutdown_tx: watch::Sender<bool>,
24    shutdown_rx: Mutex<watch::Receiver<bool>>,
25    env_provider: Arc<MockEnvProvider>,
26}
27
28impl TestDaemonContext {
29    pub fn builder(daemon_name: impl Into<String>) -> TestDaemonContextBuilder {
30        TestDaemonContextBuilder::new(daemon_name)
31    }
32
33    pub fn db(&self) -> Option<&PgPool> {
34        self.pool.as_ref()
35    }
36
37    pub fn http(&self) -> &MockHttp {
38        &self.http
39    }
40
41    pub fn is_shutdown_requested(&self) -> bool {
42        self.shutdown_rx
43            .try_lock()
44            .map(|rx| *rx.borrow())
45            .unwrap_or(false)
46    }
47
48    pub fn request_shutdown(&self) {
49        let _ = self.shutdown_tx.send(true);
50    }
51
52    /// Waits until `shutdown_tx` sends `true`.
53    pub async fn shutdown_signal(&self) {
54        let mut rx = self.shutdown_rx.lock().await;
55        while !*rx.borrow_and_update() {
56            if rx.changed().await.is_err() {
57                break;
58            }
59        }
60    }
61
62    pub async fn heartbeat(&self) -> Result<()> {
63        Ok(())
64    }
65
66    pub fn env_mock(&self) -> &MockEnvProvider {
67        &self.env_provider
68    }
69}
70
71impl EnvAccess for TestDaemonContext {
72    fn env_provider(&self) -> &dyn EnvProvider {
73        self.env_provider.as_ref()
74    }
75}
76
77pub struct TestDaemonContextBuilder {
78    daemon_name: String,
79    instance_id: Option<Uuid>,
80    pool: Option<PgPool>,
81    http: MockHttp,
82    env_vars: HashMap<String, String>,
83}
84
85impl TestDaemonContextBuilder {
86    pub fn new(daemon_name: impl Into<String>) -> Self {
87        Self {
88            daemon_name: daemon_name.into(),
89            instance_id: None,
90            pool: None,
91            http: MockHttp::new(),
92            env_vars: HashMap::new(),
93        }
94    }
95
96    pub fn with_instance_id(mut self, id: Uuid) -> Self {
97        self.instance_id = Some(id);
98        self
99    }
100
101    pub fn with_pool(mut self, pool: PgPool) -> Self {
102        self.pool = Some(pool);
103        self
104    }
105
106    pub fn mock_http<F>(self, pattern: &str, handler: F) -> Self
107    where
108        F: Fn(&MockRequest) -> MockResponse + Send + Sync + 'static,
109    {
110        self.http.add_mock_sync(pattern, handler);
111        self
112    }
113
114    pub fn mock_http_json<T: serde::Serialize>(self, pattern: &str, response: T) -> Self {
115        let json = serde_json::to_value(response).unwrap_or(serde_json::Value::Null);
116        self.mock_http(pattern, move |_| MockResponse::json(json.clone()))
117    }
118
119    pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
120        self.env_vars.insert(key.into(), value.into());
121        self
122    }
123
124    pub fn with_envs(mut self, vars: HashMap<String, String>) -> Self {
125        self.env_vars.extend(vars);
126        self
127    }
128
129    pub fn build(self) -> TestDaemonContext {
130        let (shutdown_tx, shutdown_rx) = watch::channel(false);
131
132        TestDaemonContext {
133            daemon_name: self.daemon_name,
134            instance_id: self.instance_id.unwrap_or_else(Uuid::new_v4),
135            pool: self.pool,
136            http: Arc::new(self.http),
137            shutdown_tx,
138            shutdown_rx: Mutex::new(shutdown_rx),
139            env_provider: Arc::new(MockEnvProvider::with_vars(self.env_vars)),
140        }
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn test_daemon_context_creation() {
150        let ctx = TestDaemonContext::builder("heartbeat_daemon").build();
151
152        assert_eq!(ctx.daemon_name, "heartbeat_daemon");
153        assert!(!ctx.is_shutdown_requested());
154    }
155
156    #[test]
157    fn test_shutdown_request() {
158        let ctx = TestDaemonContext::builder("test").build();
159
160        assert!(!ctx.is_shutdown_requested());
161        ctx.request_shutdown();
162        assert!(ctx.is_shutdown_requested());
163    }
164
165    #[tokio::test]
166    async fn test_shutdown_signal() {
167        let ctx = TestDaemonContext::builder("test").build();
168
169        let shutdown_tx = ctx.shutdown_tx.clone();
170        tokio::spawn(async move {
171            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
172            let _ = shutdown_tx.send(true);
173        });
174
175        tokio::time::timeout(std::time::Duration::from_millis(200), ctx.shutdown_signal())
176            .await
177            .expect("Shutdown signal should complete");
178
179        assert!(ctx.is_shutdown_requested());
180    }
181
182    #[test]
183    fn test_with_instance_id() {
184        let id = Uuid::new_v4();
185        let ctx = TestDaemonContext::builder("test")
186            .with_instance_id(id)
187            .build();
188
189        assert_eq!(ctx.instance_id, id);
190    }
191}
forge_core/testing/context/daemon.rs

forge_core/testing/context/
daemon.rs
