Add address index tests

PIVX-Labs · Sep 20, 2024 · 7fe5f99 · 7fe5f99
1 parent 42c874a
commit 7fe5f99
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 0 deletions.
diff --git a/src-tauri/src/address_index/block_source.rs b/src-tauri/src/address_index/block_source.rs
@@ -7,3 +7,36 @@ pub trait BlockSource {
         &mut self,
     ) -> crate::error::Result<Pin<Box<dyn Stream<Item = Block> + '_ + Send>>>;
 }
+
+#[cfg(test)]
+pub mod test {
+    use super::super::types::{Block, Tx};
+    use super::*;
+    pub struct MockBlockSource;
+    impl BlockSource for MockBlockSource {
+        fn get_blocks(
+            &mut self,
+        ) -> crate::error::Result<Pin<Box<dyn Stream<Item = Block> + '_ + Send>>> {
+            Ok(Box::pin(futures::stream::iter(vec![
+                Block {
+                    txs: vec![Tx {
+                        txid: "txid1".to_owned(),
+                        addresses: vec!["address1".to_owned(), "address2".to_owned()],
+                    }],
+                },
+                Block {
+                    txs: vec![Tx {
+                        txid: "txid2".to_owned(),
+                        addresses: vec!["address1".to_owned(), "address4".to_owned()],
+                    }],
+                },
+                Block {
+                    txs: vec![Tx {
+                        txid: "txid3".to_owned(),
+                        addresses: vec!["address1".to_owned(), "address5".to_owned()],
+                    }],
+                },
+            ])))
+        }
+    }
+}
diff --git a/src-tauri/src/address_index/database.rs b/src-tauri/src/address_index/database.rs
@@ -16,3 +16,30 @@ pub trait Database {
         Ok(())
     }
 }
+
+#[cfg(test)]
+pub mod test {
+    use super::*;
+    use std::collections::HashMap;
+
+    #[derive(Default)]
+    pub struct MockDB {
+        address_map: HashMap<String, Vec<String>>,
+    }
+
+    impl Database for MockDB {
+        async fn get_address_txids(&self, address: &str) -> crate::error::Result<Vec<String>> {
+            Ok(self.address_map.get(address).unwrap_or(&vec![]).clone())
+        }
+
+        async fn store_tx(&mut self, tx: &Tx) -> crate::error::Result<()> {
+            for address in &tx.addresses {
+                self.address_map
+                    .entry(address.clone())
+                    .and_modify(|vec| vec.push(tx.txid.clone()))
+                    .or_insert(vec![tx.txid.clone()]);
+            }
+            Ok(())
+        }
+    }
+}
diff --git a/src-tauri/src/address_index/mod.rs b/src-tauri/src/address_index/mod.rs
@@ -30,4 +30,43 @@ impl<D: Database + Send, B: BlockSource + Send> AddressIndex<D, B> {
             block_source,
         }
     }
+    async fn get_address_txids(&self, address: &str) -> crate::error::Result<Vec<String>> {
+        self.database.get_address_txids(address).await
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::block_source::test::MockBlockSource;
+    use super::database::test::MockDB;
+    use super::*;
+
+    #[tokio::test]
+    async fn syncs_correctly() -> crate::error::Result<()> {
+        let mock_db = MockDB::default();
+        let block_source = MockBlockSource;
+        let mut address_index = AddressIndex::new(mock_db, block_source);
+        address_index.sync().await?;
+        assert_eq!(
+            address_index.get_address_txids("address1").await?,
+            vec!["txid1", "txid2", "txid3"]
+        );
+        assert_eq!(
+            address_index.get_address_txids("address2").await?,
+            vec!["txid1"]
+        );
+        assert_eq!(
+            address_index.get_address_txids("address4").await?,
+            vec!["txid2"]
+        );
+        assert_eq!(
+            address_index.get_address_txids("address5").await?,
+            vec!["txid3"]
+        );
+        assert_eq!(
+            address_index.get_address_txids("address6").await?,
+            Vec::<String>::new()
+        );
+        Ok(())
+    }
 }