add tests

tests/test_mock_inheritance.py

@@ -0,0 +1,281 @@
+"""
+Mock tests for inheritance core logic.
+
+Tests:
+  1. prepare_transactions builds outputs matching heir amounts
+  2. prepare_transactions handles multiple locktime groups
+  3. new will item gets ANTICIPATED status via set_anticipate
+  4. REPLACED status propagates through the will tree
+  5. INVALIDATED status when input is spent and parent tx not tracked
+"""
+
+import sys, os, time, copy
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), os.pardir))
+
+from bal.core.heirs import (
+    Heirs, prepare_transactions, HEIR_ADDRESS, HEIR_REAL_AMOUNT,
+)
+from bal.core.will import Will, WillItem
+from bal.core.util import Util
+from electrum.transaction import (
+    PartialTransaction, PartialTxInput, PartialTxOutput, TxOutpoint,
+)
+from electrum.util import bfh
+
+ADDR1 = "1A1zP1eP5QGefi2DMPTfTL5SLmv7DivfNa"
+ADDR2 = "1BvBMSEYstWetqTFn5Au4m4GFg7xJaNVN2"
+ADDR3 = "12c6DSiU4Rq3P4ZxziKxzrL5LmMBrzjrJX"
+ADDR_CHG = "3J98t1WpEZ73CNmQviecrnyiWrnqRhWNLy"
+TX_FEES = 100
+
+
+def _utxo(txid_hex, value=500_000):
+    txin = PartialTxInput(prevout=TxOutpoint(txid=bfh(txid_hex), out_idx=0))
+    txin._trusted_value_sats = value
+    txin._TxInput__value_sats = value
+    txin.is_mine = True
+    return txin
+
+
+def _tx(utxos, outputs, locktime):
+    return PartialTransaction.from_io(utxos, outputs, locktime=locktime, version=2)
+
+
+def _wi(tx, heirs_data):
+    wi = WillItem.__new__(WillItem)
+    wi.tx = tx
+    wi._id = tx.txid()
+    wi.heirs = heirs_data
+    wi.we = None
+    wi.status = ""
+    wi.STATUS = copy.deepcopy(WillItem.STATUS_DEFAULT)
+    wi.tx_fees = TX_FEES
+    wi.father = None
+    wi.children = {}
+    wi.set_status("VALID", True)
+    return wi
+
+
+def _txid_gen():
+    i = 0
+    while True:
+        yield f"{i:064x}"
+        i += 1
+
+
+def _wallet():
+    w = MagicMock()
+    w.dust_threshold.return_value = 500
+    w.get_change_addresses_for_new_transaction.return_value = [ADDR_CHG]
+    w.get_utxos.return_value = []
+    w.db.get.return_value = {}
+    w.db.get_transaction.return_value = None
+    return w
+
+
+# ------------------------------------------------------------------ #
+# Test 1: prepare_transactions builds outputs matching heir amounts
+# ------------------------------------------------------------------ #
+
+@patch("electrum.transaction.bitcoin.address_to_script", return_value=b"\x00\x14" + b"\x00" * 20)
+@patch("bal.core.heirs.bitcoin.is_address", return_value=True)
+@patch("electrum.transaction.bitcoin.is_address", return_value=True)
+@patch("bal.core.heirs.PartialTransaction.txid", side_effect=_txid_gen())
+def test_prepare_transactions_verify_amounts(*_):
+    wallet = _wallet()
+    now = int(time.time())
+    lt = now + 30 * 86400
+
+    locktimes = {
+        lt: {
+            "alice":   [ADDR1, "fix", str(lt), 1_000_000],
+            "bob":     [ADDR2, "fix", str(lt), 1_000_000],
+            "charlie": [ADDR3, "fix", str(lt), 1_000_000],
+        }
+    }
+    utxos = [_utxo(f"{i:02x}" * 32, 500_000) for i in range(8)]
+    fees = {lt: 10000}
+
+    txs = prepare_transactions(locktimes, utxos[:], fees, wallet)
+    assert txs, "no txs built"
+
+    for txid, tx in txs.items():
+        outs = tx.outputs()
+        for name, hdata in getattr(tx, "heirs", {}).items():
+            if "w!ll3x3c" in name:
+                continue
+            addr = hdata[HEIR_ADDRESS]
+            exp = hdata[HEIR_REAL_AMOUNT]
+            assert any(o.value == exp for o in outs), \
+                f"{name}: amount {exp} not in outputs"
+
+    print(f"[OK] prepare_transactions built {len(txs)} tx(s)")
+
+
+# ------------------------------------------------------------------ #
+# Test 2: prepare_transactions with multiple locktime groups
+# ------------------------------------------------------------------ #
+
+@patch("electrum.transaction.bitcoin.address_to_script", return_value=b"\x00\x14" + b"\x00" * 20)
+@patch("bal.core.heirs.bitcoin.is_address", return_value=True)
+@patch("electrum.transaction.bitcoin.is_address", return_value=True)
+@patch("bal.core.heirs.PartialTransaction.txid", side_effect=_txid_gen())
+def test_prepare_transactions_multiple_calls(*_):
+    wallet = _wallet()
+    now = int(time.time())
+    lt1, lt2 = now + 30 * 86400, now + 60 * 86400
+
+    utxos_pool = [_utxo(f"{i:02x}" * 32, 600_000) for i in range(2)]
+
+    # First call: process lt1 group
+    txs1 = prepare_transactions(
+        {lt1: {"alice": [ADDR1, 500_000, str(lt1)]}},
+        utxos_pool[:], {lt1: 5000}, wallet,
+    )
+    assert txs1, "no txs from first group"
+    t1_heirs = set()
+    for tx in txs1.values():
+        t1_heirs.update(getattr(tx, "heirs", {}).keys())
+    assert "alice" in t1_heirs, "alice missing from first call"
+
+    # Second call: process lt2 group (with fresh mock data)
+    txs2 = prepare_transactions(
+        {lt2: {"bob": [ADDR2, 300_000, str(lt2)]}},
+        [_utxo("ff" * 32, 600_000)], {lt2: 5000}, wallet,
+    )
+    assert txs2, "no txs from second group"
+    t2_heirs = set()
+    for tx in txs2.values():
+        t2_heirs.update(getattr(tx, "heirs", {}).keys())
+    assert "bob" in t2_heirs, "bob missing from second call"
+
+    print("[OK] prepare_transactions handles multiple sequential calls")
+
+
+# ------------------------------------------------------------------ #
+# Test 3: set_anticipate gives ANTICIPATED status + lower locktime
+# ------------------------------------------------------------------ #
+
+def test_set_anticipate():
+    now = int(time.time())
+    lt = int((now + 30 * 86400) / 86400) * 86400
+
+    old_tx = _tx([_utxo("aa" * 32, 200_000)],
+                 [PartialTxOutput.from_address_and_value(ADDR1, 100_000)],
+                 lt)
+    old_wi = _wi(old_tx, {"alice": [ADDR1, 100_000, str(lt)]})
+
+    new_tx = _tx([_utxo("bb" * 32, 200_000)],
+                 [PartialTxOutput.from_address_and_value(ADDR2, 100_000)],
+                 lt)
+    new_wi = _wi(new_tx, {"alice": [ADDR2, 100_000, str(lt)]})
+
+    anticipated = new_wi.set_anticipate(old_wi)
+
+    assert anticipated, "set_anticipate should return True"
+    assert new_wi.get_status("ANTICIPATED"), "ANTICIPATED should be set"
+    assert new_wi.get_status("VALID"), "VALID should remain True"
+    assert new_wi.tx.locktime <= old_wi.tx.locktime, \
+        f"new locktime {new_wi.tx.locktime} should be <= old {old_wi.tx.locktime}"
+
+    print("[OK] set_anticipate sets ANTICIPATED, keeps VALID, lowers locktime")
+
+
+# ------------------------------------------------------------------ #
+# Test 4: REPLACED propagates through will tree
+# ------------------------------------------------------------------ #
+
+def test_replaced_propagates():
+    now = int(time.time())
+    lt_parent  = int(now + 60 * 86400)
+    lt_child   = int(now + 30 * 86400)
+    lt_replacer = int(now + 15 * 86400)
+
+    # Pre-determine parent txid (can't use real tx.txid() on unsigned PSBT)
+    parent_txid_hex = "aa" * 32
+
+    # Build parent tx spending the wallet utxo
+    parent_utxo = _utxo("aa" * 32, 200_000)
+    parent_tx = _tx([parent_utxo],
+                    [PartialTxOutput.from_address_and_value(ADDR1, 100_000),
+                     PartialTxOutput.from_address_and_value(ADDR_CHG, 99_000)],
+                    lt_parent)
+
+    # Build child tx spending the parent's change output (parent_txid:1)
+    ci = PartialTxInput(prevout=TxOutpoint(txid=bfh(parent_txid_hex), out_idx=1))
+    ci._trusted_value_sats = 99_000; ci.is_mine = True
+    child_tx = _tx([ci],
+                   [PartialTxOutput.from_address_and_value(ADDR2, 50_000)], lt_child)
+
+    # Build replacer tx spending the SAME utxo as parent, with lower locktime
+    replacer_tx = _tx([_utxo("aa" * 32, 200_000)],
+                      [PartialTxOutput.from_address_and_value(ADDR3, 150_000)], lt_replacer)
+
+    hp = {"a": [ADDR1, 100_000, str(lt_parent)]}
+    hc = {"b": [ADDR2,  50_000, str(lt_child)]}
+    hr = {"c": [ADDR3, 150_000, str(lt_replacer)]}
+
+    pwi = _wi(parent_tx, hp)
+    pwi._id = parent_txid_hex
+    cwi = _wi(child_tx, hc)
+    cwi._id = "bb" * 32
+    rwi = _wi(replacer_tx, hr)
+    rwi._id = "cc" * 32
+    will = {pwi._id: pwi, cwi._id: cwi, rwi._id: rwi}
+
+    Will.add_willtree(will)
+
+    all_utxos = [_utxo("aa" * 32, 200_000)]
+    Will.search_rai(Will.get_all_inputs(will, only_valid=True), all_utxos, will, _wallet())
+
+    assert pwi.get_status("REPLACED"), "parent should be REPLACED"
+    assert not pwi.get_status("VALID"), "parent should NOT be VALID"
+
+    assert cwi.get_status("REPLACED"), "child should be REPLACED (propagated)"
+    assert not cwi.get_status("VALID"), "child should NOT be VALID"
+
+    assert rwi.get_status("VALID"), "replacer should be VALID"
+    assert not rwi.get_status("REPLACED"), "replacer should NOT be REPLACED"
+
+    print("[OK] REPLACED propagation to children verified")
+
+
+# ------------------------------------------------------------------ #
+# Test 5: INVALIDATED when input not in utxos and parent TX untracked
+# ------------------------------------------------------------------ #
+
+def test_invalidated():
+    now = int(time.time())
+    lt = int(now + 30 * 86400)
+    addr = ADDR1
+
+    # TX that spends an unknown input (not in wallet utxos)
+    txin = PartialTxInput(prevout=TxOutpoint(txid=bfh("dd" * 32), out_idx=0))
+    txin._trusted_value_sats = 100_000; txin.is_mine = True
+    tx = _tx([txin],
+             [PartialTxOutput.from_address_and_value(addr, 90_000)], lt)
+    wi = _wi(tx, {"a": [addr, 90_000, str(lt)]})
+
+    will = {wi._id: wi}
+    Will.add_willtree(will)
+
+    # all_utxos is empty, wallet.db.get_transaction returns None
+    Will.search_rai(Will.get_all_inputs(will, only_valid=True), [], will, _wallet())
+
+    assert wi.get_status("INVALIDATED"), "should be INVALIDATED"
+    assert not wi.get_status("VALID"), "should NOT be VALID"
+
+    print("[OK] INVALIDATED when input not ours and parent TX untracked")
+
+
+# ------------------------------------------------------------------ #
+
+if __name__ == "__main__":
+    import logging
+    logging.basicConfig(level=logging.WARNING)
+    for name in sorted(dir()):
+        if name.startswith("test_"):
+            globals()[name]()
+    print("[OK] All mock inheritance tests passed")