Skip to main content

ops_rs/
contexts.rs

1use crate::prelude::*;
2use std::collections::HashMap;
3
4/// Control flow flags for ops execution
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct ControlFlags {
7    pub aborted: bool,
8    pub abort_reason: Option<String>,
9}
10
11impl Default for ControlFlags {
12    fn default() -> Self {
13        Self {
14            aborted: false,
15            abort_reason: None,
16        }
17    }
18}
19
20/// DryContext contains only serializable data values
21#[derive(Debug, Clone, Default, Serialize, Deserialize)]
22pub struct DryContext {
23    values: HashMap<String, serde_json::Value>,
24    control_flags: ControlFlags,
25}
26
27impl DryContext {
28    pub fn new() -> Self {
29        Self {
30            values: HashMap::new(),
31            control_flags: ControlFlags::default(),
32        }
33    }
34
35    pub fn with_value<T: Serialize>(mut self, key: impl Into<String>, value: T) -> Self {
36        self.insert(key, value);
37        self
38    }
39
40    pub fn insert<T: Serialize>(&mut self, key: impl Into<String>, value: T) {
41        self.values.insert(
42            key.into(),
43            serde_json::to_value(value).expect("Failed to serialize value"),
44        );
45    }
46
47    pub fn get<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Option<T> {
48        self.values
49            .get(key)
50            .and_then(|v| serde_json::from_value(v.clone()).ok())
51    }
52
53    pub fn get_required<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T, OpError> {
54        match self.values.get(key) {
55            None => Err(OpError::Context(format!(
56                "Required dry context key '{}' not found",
57                key
58            ))),
59            Some(value) => match serde_json::from_value::<T>(value.clone()) {
60                Ok(parsed) => Ok(parsed),
61                Err(_) => {
62                    let actual_type = match value {
63                        serde_json::Value::Null => "null",
64                        serde_json::Value::Bool(_) => "boolean",
65                        serde_json::Value::Number(_) => "number",
66                        serde_json::Value::String(_) => "string",
67                        serde_json::Value::Array(_) => "array",
68                        serde_json::Value::Object(_) => "object",
69                    };
70                    let expected_type = std::any::type_name::<T>();
71                    Err(OpError::Context(format!(
72                            "Type mismatch for dry context key '{}': expected type '{}', but found '{}' value: {}",
73                            key, expected_type, actual_type, value
74                        )))
75                }
76            },
77        }
78    }
79
80    pub fn contains(&self, key: &str) -> bool {
81        self.values.contains_key(key)
82    }
83
84    pub fn keys(&self) -> impl Iterator<Item = &String> {
85        self.values.keys()
86    }
87
88    pub fn values(&self) -> &HashMap<String, serde_json::Value> {
89        &self.values
90    }
91
92    /// Get a value or insert it using a factory closure if it doesn't exist
93    pub fn get_or_insert_with<T, F>(&mut self, key: &str, factory: F) -> Result<T, OpError>
94    where
95        T: Serialize + for<'de> Deserialize<'de>,
96        F: FnOnce() -> T,
97    {
98        if let Some(value) = self.get::<T>(key) {
99            Ok(value)
100        } else {
101            let new_value = factory();
102            self.insert(key, &new_value);
103            Ok(new_value)
104        }
105    }
106
107    /// Get a value or compute it using a closure that has access to the context
108    pub fn get_or_compute_with<T, F>(&mut self, key: &str, computer: F) -> Result<T, OpError>
109    where
110        T: Serialize + for<'de> Deserialize<'de>,
111        F: FnOnce(&mut Self, &str) -> T,
112    {
113        if let Some(value) = self.get::<T>(key) {
114            Ok(value)
115        } else {
116            let new_value = computer(self, key);
117            self.insert(key, &new_value);
118            Ok(new_value)
119        }
120    }
121
122    /// Get a value or compute it using a closure that has access to the context
123    /// The closure receives mutable access to the context and the key, and must insert the value itself
124    pub async fn ensure<T, F>(
125        &mut self,
126        key: &str,
127        wet: &mut WetContext,
128        factory: F,
129    ) -> Result<T, OpError>
130    where
131        T: Serialize + for<'de> Deserialize<'de>,
132        F: for<'a> FnOnce(
133            &'a mut Self,
134            &'a mut WetContext,
135            &'a str,
136        ) -> std::pin::Pin<
137            Box<dyn std::future::Future<Output = Result<T, OpError>> + Send + 'a>,
138        >,
139    {
140        if let Some(value) = self.get::<T>(key) {
141            Ok(value)
142        } else {
143            let new_value = factory(self, wet, key).await?;
144            self.insert(key, &new_value);
145
146            Ok(new_value)
147        }
148    }
149
150    pub fn merge(&mut self, other: DryContext) {
151        self.values.extend(other.values);
152        // Only merge control flags if they are set in other and not already set in self
153        if other.control_flags.aborted && !self.control_flags.aborted {
154            self.control_flags.aborted = true;
155            self.control_flags.abort_reason = other.control_flags.abort_reason;
156        }
157    }
158
159    /// Set abort flag with optional reason
160    pub fn set_abort(&mut self, reason: Option<String>) {
161        self.control_flags.aborted = true;
162        self.control_flags.abort_reason = reason;
163    }
164
165    /// Check if abort flag is set
166    pub fn is_aborted(&self) -> bool {
167        self.control_flags.aborted
168    }
169
170    /// Get abort reason if set
171    pub fn abort_reason(&self) -> Option<&String> {
172        self.control_flags.abort_reason.as_ref()
173    }
174
175    /// Clear all control flags
176    pub fn clear_control_flags(&mut self) {
177        self.control_flags = ControlFlags::default();
178    }
179}
180
181/// WetContext contains runtime references (services, connections, etc.)
182#[derive(Debug, Default)]
183pub struct WetContext {
184    references: HashMap<String, Arc<dyn Any + Send + Sync>>,
185}
186
187// WetContext is Send and Sync because all its contents are Send + Sync
188unsafe impl Send for WetContext {}
189unsafe impl Sync for WetContext {}
190
191impl WetContext {
192    pub fn new() -> Self {
193        Self {
194            references: HashMap::new(),
195        }
196    }
197
198    pub fn with_ref<T: Any + Send + Sync>(mut self, key: impl Into<String>, value: T) -> Self {
199        self.insert_ref(key, value);
200        self
201    }
202
203    pub fn insert_ref<T: Any + Send + Sync>(&mut self, key: impl Into<String>, value: T) {
204        self.references.insert(key.into(), Arc::new(value));
205    }
206
207    pub fn insert_arc(&mut self, key: impl Into<String>, value: Arc<dyn Any + Send + Sync>) {
208        self.references.insert(key.into(), value);
209    }
210
211    pub fn get_ref<T: Any + Send + Sync>(&self, key: &str) -> Option<Arc<T>> {
212        self.references
213            .get(key)
214            .and_then(|any_ref| any_ref.clone().downcast::<T>().ok())
215    }
216
217    pub fn get_required<T: Any + Send + Sync>(&self, key: &str) -> Result<Arc<T>, OpError> {
218        match self.references.get(key) {
219            None => Err(OpError::Context(format!(
220                "Required wet context reference '{}' not found",
221                key
222            ))),
223            Some(any_ref) => match any_ref.clone().downcast::<T>() {
224                Ok(typed_ref) => Ok(typed_ref),
225                Err(_) => {
226                    let expected_type = std::any::type_name::<T>();
227                    Err(OpError::Context(format!(
228                            "Type mismatch for wet context reference '{}': expected type '{}', but found a different type",
229                            key, expected_type
230                        )))
231                }
232            },
233        }
234    }
235
236    pub fn contains(&self, key: &str) -> bool {
237        self.references.contains_key(key)
238    }
239
240    pub fn keys(&self) -> impl Iterator<Item = &String> {
241        self.references.keys()
242    }
243
244    /// Get a reference or compute it using an async closure that has access to both contexts
245    pub async fn ensure<T, F>(
246        &mut self,
247        key: &str,
248        dry: &mut DryContext,
249        factory: F,
250    ) -> Result<Arc<T>, OpError>
251    where
252        T: Any + Send + Sync,
253        F: for<'a> FnOnce(
254            &'a mut DryContext,
255            &'a mut Self,
256            &'a str,
257        ) -> std::pin::Pin<
258            Box<dyn std::future::Future<Output = Result<Arc<T>, OpError>> + Send + 'a>,
259        >,
260    {
261        if let Some(value) = self.get_ref::<T>(key) {
262            Ok(value)
263        } else {
264            let new_value = factory(dry, self, key).await?;
265            self.insert_arc(key, new_value.clone());
266            Ok(new_value)
267        }
268    }
269
270    pub fn merge(&mut self, other: WetContext) {
271        self.references.extend(other.references);
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278
279    // TEST0009: Insert typed values into DryContext and verify get/contains work correctly
280    #[test]
281    fn test0009_dry_context_basic_operations() {
282        let mut ctx = DryContext::new();
283        ctx.insert("name", "test");
284        ctx.insert("count", 42);
285
286        assert_eq!(ctx.get::<String>("name").unwrap(), "test");
287        assert_eq!(ctx.get::<i32>("count").unwrap(), 42);
288        assert!(ctx.contains("name"));
289        assert!(!ctx.contains("missing"));
290    }
291
292    // TEST0010: Build a DryContext with chained with_value calls and verify all values are stored
293    #[test]
294    fn test0010_dry_context_builder() {
295        let ctx = DryContext::new()
296            .with_value("key1", "value1")
297            .with_value("key2", 123);
298
299        assert_eq!(ctx.get::<String>("key1").unwrap(), "value1");
300        assert_eq!(ctx.get::<i32>("key2").unwrap(), 123);
301    }
302
303    // TEST0011: Insert a reference into WetContext and retrieve it by type via get_ref
304    #[test]
305    fn test0011_wet_context_basic_operations() {
306        #[derive(Debug)]
307        struct TestService {
308            name: String,
309        }
310
311        let mut ctx = WetContext::new();
312        let service = TestService {
313            name: "test".to_string(),
314        };
315        ctx.insert_ref("service", service);
316
317        let retrieved = ctx.get_ref::<TestService>("service").unwrap();
318        assert_eq!(retrieved.name, "test");
319    }
320
321    // TEST0012: Build a WetContext with chained with_ref calls and verify contains for each key
322    #[test]
323    fn test0012_wet_context_builder() {
324        struct Service1;
325        struct Service2;
326
327        let ctx = WetContext::new()
328            .with_ref("service1", Service1)
329            .with_ref("service2", Service2);
330
331        assert!(ctx.contains("service1"));
332        assert!(ctx.contains("service2"));
333    }
334
335    // TEST0013: Confirm get_required succeeds for present keys and returns an error for missing keys
336    #[test]
337    fn test0013_required_values() {
338        let ctx = DryContext::new().with_value("exists", 42);
339
340        assert_eq!(ctx.get_required::<i32>("exists").unwrap(), 42);
341        assert!(ctx.get_required::<i32>("missing").is_err());
342    }
343
344    // TEST0014: Merge two DryContexts and verify values from both are accessible in the target
345    #[test]
346    fn test0014_context_merge() {
347        let mut ctx1 = DryContext::new().with_value("a", 1);
348        let ctx2 = DryContext::new().with_value("b", 2);
349
350        ctx1.merge(ctx2);
351        assert_eq!(ctx1.get::<i32>("a").unwrap(), 1);
352        assert_eq!(ctx1.get::<i32>("b").unwrap(), 2);
353    }
354
355    // TEST0015: Verify get_required returns a Type mismatch error when the stored type doesn't match
356    #[test]
357    fn test0015_dry_context_type_mismatch_error() {
358        let ctx = DryContext::new()
359            .with_value("count", "not_a_number")
360            .with_value("flag", 123);
361
362        // String value, expecting i32
363        let result = ctx.get_required::<i32>("count");
364        assert!(result.is_err());
365        let err = result.unwrap_err().to_string();
366        assert!(err.contains("Type mismatch"));
367        assert!(err.contains("expected type 'i32'"));
368        assert!(err.contains("found 'string' value"));
369
370        // Number value, expecting bool
371        let result = ctx.get_required::<bool>("flag");
372        assert!(result.is_err());
373        let err = result.unwrap_err().to_string();
374        assert!(err.contains("Type mismatch"));
375        assert!(err.contains("expected type 'bool'"));
376        assert!(err.contains("found 'number' value"));
377
378        // Missing key still gives "not found"
379        let result = ctx.get_required::<i32>("missing");
380        assert!(result.is_err());
381        let err = result.unwrap_err().to_string();
382        assert!(err.contains("not found"));
383        assert!(!err.contains("Type mismatch"));
384    }
385
386    // TEST0016: Verify WetContext get_required returns a Type mismatch error when the stored ref type differs
387    #[test]
388    fn test0016_wet_context_type_mismatch_error() {
389        #[derive(Debug)]
390        struct ServiceA {
391            _name: String,
392        }
393        #[derive(Debug)]
394        struct ServiceB {
395            _id: i32,
396        }
397
398        let mut ctx = WetContext::new();
399        ctx.insert_ref(
400            "service",
401            ServiceA {
402                _name: "test".to_string(),
403            },
404        );
405
406        // Wrong type
407        let result = ctx.get_required::<ServiceB>("service");
408        assert!(result.is_err());
409        let err = result.unwrap_err().to_string();
410        assert!(err.contains("Type mismatch"));
411        assert!(err.contains("expected type"));
412        assert!(err.contains("ServiceB"));
413
414        // Missing key
415        let result = ctx.get_required::<ServiceA>("missing");
416        assert!(result.is_err());
417        let err = result.unwrap_err().to_string();
418        assert!(err.contains("not found"));
419        assert!(!err.contains("Type mismatch"));
420    }
421
422    // TEST0017: Set and clear abort flags on DryContext and verify is_aborted and abort_reason reflect state
423    #[test]
424    fn test0017_control_flags() {
425        let mut ctx = DryContext::new();
426
427        // Test abort functionality
428        assert!(!ctx.is_aborted());
429        assert_eq!(ctx.abort_reason(), None);
430
431        ctx.set_abort(Some("Test abort reason".to_string()));
432        assert!(ctx.is_aborted());
433        assert_eq!(ctx.abort_reason(), Some(&"Test abort reason".to_string()));
434
435        // Test clearing all flags
436        ctx.set_abort(Some("Another reason".to_string()));
437        assert!(ctx.is_aborted());
438
439        ctx.clear_control_flags();
440        assert!(!ctx.is_aborted());
441        assert_eq!(ctx.abort_reason(), None);
442    }
443
444    // TEST0018: Merge contexts with abort flags and confirm the target inherits the abort state correctly
445    #[test]
446    fn test0018_control_flags_merge() {
447        let mut ctx1 = DryContext::new();
448        let mut ctx2 = DryContext::new();
449
450        // Set flags in ctx2
451        ctx2.set_abort(Some("Merged abort".to_string()));
452
453        // Merge ctx2 into ctx1
454        ctx1.merge(ctx2);
455
456        assert!(ctx1.is_aborted());
457        assert_eq!(ctx1.abort_reason(), Some(&"Merged abort".to_string()));
458
459        // Test that merge doesn't override existing abort
460        let mut ctx3 = DryContext::new();
461        ctx3.set_abort(Some("Original abort".to_string()));
462
463        let mut ctx4 = DryContext::new();
464        ctx4.set_abort(Some("New abort".to_string()));
465
466        ctx3.merge(ctx4);
467        // Should keep the original abort reason since ctx3 was already aborted
468        assert_eq!(ctx3.abort_reason(), Some(&"Original abort".to_string()));
469    }
470
471    // TEST0019: Verify get_or_insert_with inserts when missing and returns existing without calling factory
472    #[test]
473    fn test0019_get_or_insert_with() {
474        let mut ctx = DryContext::new();
475
476        // Test inserting a new value when key doesn't exist
477        let value = ctx.get_or_insert_with("count", || 42).unwrap();
478        assert_eq!(value, 42);
479        assert_eq!(ctx.get::<i32>("count").unwrap(), 42);
480
481        // Test getting existing value without calling factory
482        let mut factory_called = false;
483        let value = ctx
484            .get_or_insert_with("count", || {
485                factory_called = true;
486                100
487            })
488            .unwrap();
489        assert_eq!(value, 42); // Should return existing value
490        assert!(!factory_called); // Factory should not be called
491
492        // Test with different types
493        let name = ctx
494            .get_or_insert_with("name", || "default_name".to_string())
495            .unwrap();
496        assert_eq!(name, "default_name");
497        assert_eq!(ctx.get::<String>("name").unwrap(), "default_name");
498
499        // Test with complex types
500        #[derive(Debug, PartialEq, Serialize, Deserialize)]
501        struct Config {
502            host: String,
503            port: u16,
504        }
505
506        let config = ctx
507            .get_or_insert_with("config", || Config {
508                host: "localhost".to_string(),
509                port: 8080,
510            })
511            .unwrap();
512
513        assert_eq!(config.host, "localhost");
514        assert_eq!(config.port, 8080);
515
516        // Verify it's stored in context
517        let stored_config = ctx.get::<Config>("config").unwrap();
518        assert_eq!(stored_config, config);
519    }
520
521    // TEST0098: Merge two DryContexts where keys overlap and verify the merging context's values win
522    #[test]
523    fn test0098_dry_context_merge_overwrites_keys() {
524        let mut ctx1 = DryContext::new()
525            .with_value("shared", 1i32)
526            .with_value("only_in_1", 10i32);
527        let ctx2 = DryContext::new()
528            .with_value("shared", 2i32)
529            .with_value("only_in_2", 20i32);
530        ctx1.merge(ctx2);
531        // After merge, ctx2's value wins for overlapping keys
532        assert_eq!(ctx1.get::<i32>("shared").unwrap(), 2);
533        assert_eq!(ctx1.get::<i32>("only_in_1").unwrap(), 10);
534        assert_eq!(ctx1.get::<i32>("only_in_2").unwrap(), 20);
535    }
536
537    // TEST0099: Merge two WetContexts and verify both sets of references are accessible in the target
538    #[test]
539    fn test0099_wet_context_merge() {
540        struct ServiceA;
541        struct ServiceB;
542
543        let mut ctx1 = WetContext::new();
544        ctx1.insert_ref("a", ServiceA);
545
546        let mut ctx2 = WetContext::new();
547        ctx2.insert_ref("b", ServiceB);
548
549        ctx1.merge(ctx2);
550        assert!(ctx1.contains("a"));
551        assert!(ctx1.contains("b"));
552    }
553
554    // TEST0100: Serialize and deserialize a DryContext and verify all values survive the round-trip
555    #[test]
556    fn test0100_dry_context_serde_roundtrip() {
557        let original = DryContext::new()
558            .with_value("name", "alice")
559            .with_value("count", 42i32)
560            .with_value("flag", true);
561
562        let json = serde_json::to_string(&original).expect("serialize failed");
563        let restored: DryContext = serde_json::from_str(&json).expect("deserialize failed");
564
565        assert_eq!(restored.get::<String>("name").unwrap(), "alice");
566        assert_eq!(restored.get::<i32>("count").unwrap(), 42);
567        assert_eq!(restored.get::<bool>("flag").unwrap(), true);
568    }
569
570    // TEST0101: Clone a DryContext and verify the clone is independent (mutations don't propagate)
571    #[test]
572    fn test0101_dry_context_clone_is_independent() {
573        let original = DryContext::new().with_value("x", 1i32);
574        let mut cloned = original.clone();
575        cloned.insert("x", 99i32);
576        assert_eq!(original.get::<i32>("x").unwrap(), 1);
577        assert_eq!(cloned.get::<i32>("x").unwrap(), 99);
578    }
579
580    // TEST0102: Verify DryContext::keys() returns all inserted keys
581    #[test]
582    fn test0102_dry_context_keys() {
583        let ctx = DryContext::new()
584            .with_value("alpha", 1i32)
585            .with_value("beta", 2i32)
586            .with_value("gamma", 3i32);
587        let mut keys: Vec<_> = ctx.keys().cloned().collect();
588        keys.sort();
589        assert_eq!(keys, vec!["alpha", "beta", "gamma"]);
590    }
591
592    // TEST0103: Verify WetContext::keys() returns all inserted reference keys
593    #[test]
594    fn test0103_wet_context_keys() {
595        struct Svc;
596        let mut ctx = WetContext::new();
597        ctx.insert_ref("svc1", Svc);
598        ctx.insert_ref("svc2", Svc);
599        let mut keys: Vec<_> = ctx.keys().cloned().collect();
600        keys.sort();
601        assert_eq!(keys, vec!["svc1", "svc2"]);
602    }
603
604    // TEST0020: Verify get_or_compute_with computes and stores a value using context data and skips recompute if present
605    #[test]
606    fn test0020_get_or_compute_with() {
607        let mut ctx = DryContext::new();
608
609        // Seed some initial data
610        ctx.insert("base_port", 8000);
611        ctx.insert("app_name", "test_app".to_string());
612
613        // Test computing a value that depends on existing context data
614        let computed_url = ctx
615            .get_or_compute_with("service_url", |ctx, key| {
616                let base_port: i32 = ctx.get("base_port").unwrap_or(3000);
617                let app_name: String = ctx.get("app_name").unwrap_or_else(|| "default".to_string());
618                let url = format!("http://{}:{}", app_name, base_port + 80);
619
620                // The closure can insert additional related data
621                ctx.insert("computed_port", base_port + 80);
622                ctx.insert(format!("{}_timestamp", key), "2023-01-01T00:00:00Z");
623
624                url
625            })
626            .unwrap();
627
628        assert_eq!(computed_url, "http://test_app:8080");
629        assert_eq!(
630            ctx.get::<String>("service_url").unwrap(),
631            "http://test_app:8080"
632        );
633        assert_eq!(ctx.get::<i32>("computed_port").unwrap(), 8080);
634        assert_eq!(
635            ctx.get::<String>("service_url_timestamp").unwrap(),
636            "2023-01-01T00:00:00Z"
637        );
638
639        // Test getting existing value without calling computer
640        let mut computer_called = false;
641        let existing_url = ctx
642            .get_or_compute_with("service_url", |_ctx, _key| {
643                computer_called = true;
644                "should_not_be_called".to_string()
645            })
646            .unwrap();
647
648        assert_eq!(existing_url, "http://test_app:8080");
649        assert!(!computer_called);
650
651        // Test computer that doesn't insert the value (fallback insertion)
652        let fallback_value = ctx
653            .get_or_compute_with("fallback_test", |_ctx, _key| {
654                // Computer doesn't insert the value itself
655                "fallback_computed".to_string()
656            })
657            .unwrap();
658
659        assert_eq!(fallback_value, "fallback_computed");
660        assert_eq!(
661            ctx.get::<String>("fallback_test").unwrap(),
662            "fallback_computed"
663        );
664    }
665}