[bdk_chain_redesign] Add tests for IndexedTxGraph with LocalChain

These tests cover list_txout, list_utxo and balance methods.

crates/chain/tests/test_indexed_tx_graph.rs

@@ -1,12 +1,15 @@
+#[macro_use]
 mod common;
 
+use std::collections::{BTreeMap, BTreeSet};
+
 use bdk_chain::{
     indexed_tx_graph::{IndexedAdditions, IndexedTxGraph},
-    keychain::{DerivationAdditions, KeychainTxOutIndex},
+    keychain::{Balance, DerivationAdditions, KeychainTxOutIndex},
     tx_graph::Additions,
-    BlockId,
+    BlockId, ObservedAs,
 };
-use bitcoin::{secp256k1::Secp256k1, OutPoint, Transaction, TxIn, TxOut};
+use bitcoin::{secp256k1::Secp256k1, BlockHash, OutPoint, Script, Transaction, TxIn, TxOut};
 use miniscript::Descriptor;
 
 /// Ensure [`IndexedTxGraph::insert_relevant_txs`] can successfully index transactions NOT presented
@@ -71,3 +74,467 @@ fn insert_relevant_txs() {
         }
     )
 }
+
+#[test]
+fn test_list_owned_txouts() {
+    let mut local_chain = local_chain![
+        (0, h!("Block 0")),
+        (1, h!("Block 1")),
+        (2, h!("Block 2")),
+        (3, h!("Block 3"))
+    ];
+
+    let desc_1 : &str = "tr(tprv8ZgxMBicQKsPd3krDUsBAmtnRsK3rb8u5yi1zhQgMhF1tR8MW7xfE4rnrbbsrbPR52e7rKapu6ztw1jXveJSCGHEriUGZV7mCe88duLp5pj/86'/1'/0'/0/*)";
+    let (desc_1, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), desc_1).unwrap();
+    let desc_2 : &str = "tr(tprv8ZgxMBicQKsPd3krDUsBAmtnRsK3rb8u5yi1zhQgMhF1tR8MW7xfE4rnrbbsrbPR52e7rKapu6ztw1jXveJSCGHEriUGZV7mCe88duLp5pj/86'/1'/0'/1/*)";
+    let (desc_2, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), desc_2).unwrap();
+
+    let mut graph = IndexedTxGraph::<BlockId, KeychainTxOutIndex<String>>::default();
+
+    graph.index.add_keychain("keychain_1".into(), desc_1);
+    graph.index.add_keychain("keychain_2".into(), desc_2);
+
+    graph.index.set_lookahead_for_all(10);
+
+    let mut trusted_spks = Vec::new();
+    let mut untrusted_spks = Vec::new();
+
+    {
+        for _ in 0..10 {
+            let ((_, script), _) = graph.index.reveal_next_spk(&"keychain_1".to_string());
+            trusted_spks.push(script.clone());
+        }
+    }
+
+    {
+        for _ in 0..10 {
+            let ((_, script), _) = graph.index.reveal_next_spk(&"keychain_2".to_string());
+            untrusted_spks.push(script.clone());
+        }
+    }
+
+    let trust_predicate = |spk: &Script| trusted_spks.contains(spk);
+
+    // tx1 is coinbase transaction received at trusted keychain at block 0.
+    let tx1 = Transaction {
+        input: vec![TxIn {
+            previous_output: OutPoint::null(),
+            ..Default::default()
+        }],
+        output: vec![TxOut {
+            value: 70000,
+            script_pubkey: trusted_spks[0].clone(),
+        }],
+        ..common::new_tx(0)
+    };
+
+    // tx2 is an incoming transaction received at untrusted keychain at block 1.
+    let tx2 = Transaction {
+        output: vec![TxOut {
+            value: 30000,
+            script_pubkey: untrusted_spks[0].clone(),
+        }],
+        ..common::new_tx(0)
+    };
+
+    // tx3 spends tx2 and gives a change back in trusted keychain. Confirmed at Block 2.
+    let tx3 = Transaction {
+        input: vec![TxIn {
+            previous_output: OutPoint::new(tx2.txid(), 0),
+            ..Default::default()
+        }],
+        output: vec![TxOut {
+            value: 10000,
+            script_pubkey: trusted_spks[1].clone(),
+        }],
+        ..common::new_tx(0)
+    };
+
+    // tx4 is an external transaction receiving at untrusted keychain, unconfirmed.
+    let tx4 = Transaction {
+        output: vec![TxOut {
+            value: 20000,
+            script_pubkey: untrusted_spks[1].clone(),
+        }],
+        ..common::new_tx(0)
+    };
+
+    // tx5 is spending tx3 and receiving change at trusted keychain, unconfirmed.
+    let tx5 = Transaction {
+        output: vec![TxOut {
+            value: 15000,
+            script_pubkey: trusted_spks[2].clone(),
+        }],
+        ..common::new_tx(0)
+    };
+
+    // tx6 is an unrelated transaction confirmed at 3.
+    let tx6 = common::new_tx(0);
+
+    let _ = graph.insert_relevant_txs(
+        [&tx1, &tx2, &tx3, &tx6]
+            .iter()
+            .enumerate()
+            .map(|(i, tx)| (*tx, [local_chain.get_block(i as u32).unwrap()])),
+        None,
+    );
+
+    let _ = graph.insert_relevant_txs([&tx4, &tx5].iter().map(|tx| (*tx, None)), Some(100));
+
+    // AT Block 0
+    {
+        let txouts = graph
+            .list_owned_txouts(&local_chain, local_chain.get_block(0).unwrap())
+            .collect::<Vec<_>>();
+
+        let utxos = graph
+            .list_owned_unspents(&local_chain, local_chain.get_block(0).unwrap())
+            .collect::<Vec<_>>();
+
+        let balance = graph.balance(
+            &local_chain,
+            local_chain.get_block(0).unwrap(),
+            trust_predicate,
+        );
+
+        let confirmed_txouts_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_txout_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let confirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        assert_eq!(txouts.len(), 5);
+        assert_eq!(utxos.len(), 4);
+
+        assert_eq!(confirmed_txouts_txid, [tx1.txid()].into());
+        assert_eq!(
+            unconfirmed_txout_txid,
+            [tx2.txid(), tx3.txid(), tx4.txid(), tx5.txid()].into()
+        );
+
+        assert_eq!(confirmed_utxos_txid, [tx1.txid()].into());
+        assert_eq!(
+            unconfirmed_utxos_txid,
+            [tx3.txid(), tx4.txid(), tx5.txid()].into()
+        );
+
+        assert_eq!(
+            balance,
+            Balance {
+                immature: 70000,          // immature coinbase
+                trusted_pending: 25000,   // tx3 + tx5
+                untrusted_pending: 20000, // tx4
+                confirmed: 0              // Nothing is confirmed yet
+            }
+        );
+    }
+
+    // AT Block 1
+    {
+        let txouts = graph
+            .list_owned_txouts(&local_chain, local_chain.get_block(1).unwrap())
+            .collect::<Vec<_>>();
+
+        let utxos = graph
+            .list_owned_unspents(&local_chain, local_chain.get_block(1).unwrap())
+            .collect::<Vec<_>>();
+
+        let balance = graph.balance(
+            &local_chain,
+            local_chain.get_block(1).unwrap(),
+            trust_predicate,
+        );
+
+        let confirmed_txouts_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_txout_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let confirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        assert_eq!(txouts.len(), 5);
+        assert_eq!(utxos.len(), 4);
+
+        // tx2 gets into confirmed txout set
+        assert_eq!(confirmed_txouts_txid, [tx1.txid(), tx2.txid()].into());
+        assert_eq!(
+            unconfirmed_txout_txid,
+            [tx3.txid(), tx4.txid(), tx5.txid()].into()
+        );
+
+        // tx2 doesn't get into confirmed utxos set
+        assert_eq!(confirmed_utxos_txid, [tx1.txid()].into());
+        assert_eq!(
+            unconfirmed_utxos_txid,
+            [tx3.txid(), tx4.txid(), tx5.txid()].into()
+        );
+
+        // Balance breakup remains same
+        assert_eq!(
+            balance,
+            Balance {
+                immature: 70000,          // immature coinbase
+                trusted_pending: 25000,   // tx3 + tx5
+                untrusted_pending: 20000, // tx4
+                confirmed: 0              // Nothing is confirmed yet
+            }
+        );
+    }
+
+    // AT Block 2
+    {
+        let txouts = graph
+            .list_owned_txouts(&local_chain, local_chain.get_block(2).unwrap())
+            .collect::<Vec<_>>();
+
+        let utxos = graph
+            .list_owned_unspents(&local_chain, local_chain.get_block(2).unwrap())
+            .collect::<Vec<_>>();
+
+        let balance = graph.balance(
+            &local_chain,
+            local_chain.get_block(2).unwrap(),
+            trust_predicate,
+        );
+
+        let confirmed_txouts_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_txout_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let confirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        assert_eq!(txouts.len(), 5);
+        assert_eq!(utxos.len(), 4);
+
+        // tx3 now gets into the confirmed txout set
+        assert_eq!(
+            confirmed_txouts_txid,
+            [tx1.txid(), tx2.txid(), tx3.txid()].into()
+        );
+        assert_eq!(unconfirmed_txout_txid, [tx4.txid(), tx5.txid()].into());
+
+        // tx3 also gets into confirmed utxo set
+        assert_eq!(confirmed_utxos_txid, [tx1.txid(), tx3.txid()].into());
+        assert_eq!(unconfirmed_utxos_txid, [tx4.txid(), tx5.txid()].into());
+
+        assert_eq!(
+            balance,
+            Balance {
+                immature: 70000,          // immature coinbase
+                trusted_pending: 15000,   // tx5
+                untrusted_pending: 20000, // tx4
+                confirmed: 10000          // tx3 got confirmed
+            }
+        );
+    }
+
+    // AT Block 110
+    {
+        let mut local_chain_extension = (4..150)
+            .map(|i| (i as u32, h!("random")))
+            .collect::<BTreeMap<u32, BlockHash>>();
+
+        local_chain_extension.insert(3, h!("Block 3"));
+
+        local_chain
+            .apply_update(local_chain_extension.into())
+            .unwrap();
+
+        let txouts = graph
+            .list_owned_txouts(&local_chain, local_chain.get_block(110).unwrap())
+            .collect::<Vec<_>>();
+
+        let utxos = graph
+            .list_owned_unspents(&local_chain, local_chain.get_block(110).unwrap())
+            .collect::<Vec<_>>();
+
+        let balance = graph.balance(
+            &local_chain,
+            local_chain.get_block(110).unwrap(),
+            trust_predicate,
+        );
+
+        let confirmed_txouts_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_txout_txid = txouts
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let confirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Confirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        let unconfirmed_utxos_txid = utxos
+            .iter()
+            .filter_map(|full_txout| {
+                if matches!(full_txout.chain_position, ObservedAs::Unconfirmed(_)) {
+                    Some(full_txout.outpoint.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<BTreeSet<_>>();
+
+        println!("TxOuts : {:#?}", txouts);
+        println!("UTXOS {:#?}", utxos);
+        println!("{:#?}", balance);
+
+        assert_eq!(txouts.len(), 5);
+        assert_eq!(utxos.len(), 4);
+
+        assert_eq!(
+            confirmed_txouts_txid,
+            [tx1.txid(), tx2.txid(), tx3.txid()].into()
+        );
+        assert_eq!(unconfirmed_txout_txid, [tx4.txid(), tx5.txid()].into());
+
+        assert_eq!(confirmed_utxos_txid, [tx1.txid(), tx3.txid()].into());
+        assert_eq!(unconfirmed_utxos_txid, [tx4.txid(), tx5.txid()].into());
+
+        assert_eq!(
+            balance,
+            Balance {
+                immature: 0,              // immature coinbase
+                trusted_pending: 15000,   // tx5
+                untrusted_pending: 20000, // tx4
+                confirmed: 80000          // tx1 got matured
+            }
+        );
+    }
+}

crates/chain/tests/test_keychain_txout_index.rs

@@ -293,7 +293,6 @@ fn test_wildcard_derivations() {
     let _ = txout_index.reveal_to_target(&TestKeychain::External, 25);
 
     (0..=15)
-        .into_iter()
         .chain(vec![17, 20, 23].into_iter())
         .for_each(|index| assert!(txout_index.mark_used(&TestKeychain::External, index)));
 
@@ -310,7 +309,7 @@ fn test_wildcard_derivations() {
 
     // - Use all the derived till 26.
     // - next_unused() = ((27, <spk>), DerivationAdditions)
-    (0..=26).into_iter().for_each(|index| {
+    (0..=26).for_each(|index| {
         txout_index.mark_used(&TestKeychain::External, index);
     });