feat: add CapsuleReader mocking

rearch/src/capsule_reader.rs

@@ -1,18 +1,26 @@
-use std::any::TypeId;
+use std::{
+    any::{Any, TypeId},
+    collections::HashMap,
+};
 
-use crate::{Capsule, ContainerWriteTxn};
+use crate::{Capsule, CapsuleData, ContainerWriteTxn};
 
 /// Allows you to read the current data of capsules based on the given state of the container txn.
-// TODO switch this to an Enum of Actual/Mock and don't need mock feature
-// TODO mock utility, like MockCapsuleReaderBuilder::new().set(capsule, value).set(...).build()
-pub struct CapsuleReader<'scope, 'total> {
-    id: TypeId,
-    txn: &'scope mut ContainerWriteTxn<'total>,
+pub enum CapsuleReader<'scope, 'total> {
+    // For normal operation
+    Normal {
+        id: TypeId,
+        txn: &'scope mut ContainerWriteTxn<'total>,
+    },
+    // To enable easy mocking in testing
+    Mock {
+        mocks: HashMap<TypeId, Box<dyn CapsuleData>>,
+    },
 }
 
 impl<'scope, 'total> CapsuleReader<'scope, 'total> {
     pub(crate) fn new(id: TypeId, txn: &'scope mut ContainerWriteTxn<'total>) -> Self {
-        Self { id, txn }
+        Self::Normal { id, txn }
     }
 
     /// Reads the current data of the supplied capsule, initializing it if needed.
@@ -22,27 +30,46 @@ impl<'scope, 'total> CapsuleReader<'scope, 'total> {
     /// # Panics
     /// Panics when a capsule attempts to read itself in its first build.
     pub fn read<C: Capsule>(&mut self, capsule: C) -> C::Data {
-        let (this, other) = (self.id, TypeId::of::<C>());
-        if this == other {
-            return self.txn.try_read(&capsule).unwrap_or_else(|| {
-                let capsule_name = std::any::type_name::<C>();
-                panic!(
-                    "Capsule {capsule_name} tried to read itself on its first build! {} {} {}",
-                    "This is disallowed since the capsule doesn't have any data to read yet.",
-                    "To avoid this issue, wrap the `read({capsule_name})` call in an if statement",
-                    "with the `IsFirstBuildEffect`."
-                );
-            });
-        }
+        match self {
+            CapsuleReader::Normal { id, txn } => {
+                let (this, other) = (*id, TypeId::of::<C>());
+                if this == other {
+                    return txn.try_read(&capsule).unwrap_or_else(|| {
+                        let name = std::any::type_name::<C>();
+                        panic!(
+                            "Capsule {name} tried to read itself on its first build! {} {} {}",
+                            "This is disallowed since the capsule doesn't have data to read yet.",
+                            "To avoid this issue, wrap the `read({name})` call in an if statement",
+                            "with the builtin \"is first build\" side effect."
+                        );
+                    });
+                }
 
-        // Get the value (and make sure the other manager is initialized!)
-        let data = self.txn.read_or_init(capsule);
+                // Get the value (and make sure the other manager is initialized!)
+                let data = txn.read_or_init(capsule);
 
-        // Take care of some dependency housekeeping
-        self.txn.node_or_panic(other).dependents.insert(this);
-        self.txn.node_or_panic(this).dependencies.insert(other);
+                // Take care of some dependency housekeeping
+                txn.node_or_panic(other).dependents.insert(this);
+                txn.node_or_panic(this).dependencies.insert(other);
 
-        data
+                data
+            }
+            CapsuleReader::Mock { mocks } => {
+                let id = TypeId::of::<C>();
+                let any: Box<dyn Any> = mocks
+                    .get(&id)
+                    .unwrap_or_else(|| {
+                        panic!(
+                            "Mock CapsuleReader was used to read {} {}",
+                            std::any::type_name::<C>(),
+                            "when it was not included in the mock!"
+                        );
+                    })
+                    .clone();
+                *any.downcast::<C::Data>()
+                    .expect("Types should be properly enforced due to generics")
+            }
+        }
     }
 }
 
@@ -58,3 +85,22 @@ impl<A: Capsule> FnMut<(A,)> for CapsuleReader<'_, '_> {
         self.read(args.0)
     }
 }
+
+#[derive(Clone, Default)]
+pub struct MockCapsuleReaderBuilder(HashMap<TypeId, Box<dyn CapsuleData>>);
+
+impl MockCapsuleReaderBuilder {
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn set<C: Capsule>(mut self, _capsule: &C, data: C::Data) -> Self {
+        self.0.insert(TypeId::of::<C>(), Box::new(data));
+        self
+    }
+
+    pub fn build(self) -> CapsuleReader<'static, 'static> {
+        CapsuleReader::Mock { mocks: self.0 }
+    }
+}