git / moonbite-network / tests/test_wallet.py · raw

"""Tests for MyCoin Milestone 8: Wallet & Privacy."""import sysfrom pathlib import Pathsys.path.insert(0, str(Path(__file__).resolve().parent.parent))import osimport pytestfrom transaction import Transaction, TxInput, TxOutputfrom wallet import (    Wallet,    address_from_pubkey_hash,    base58_decode,    base58_encode,    is_valid_address,    pubkey_hash_from_address,)# --------------------------------------------------------------------------- ## In-memory UTXO helpers# --------------------------------------------------------------------------- #def make_utxos_paying(pkh_to_amount, prefix="fund"):    """Create a funding Transaction paying each (pkh, amount) and return its    UTXOs as a list of (txid, index, TxOutput)."""    outputs = [TxOutput(amount, pkh) for pkh, amount in pkh_to_amount]    tx = Transaction(inputs=[], outputs=outputs)    txid = tx.txid    return [(txid, i, out) for i, out in enumerate(outputs)]def make_resolver(utxos):    """Build a resolver mapping (txid, index) -> TxOutput for Transaction.verify."""    lookup = {(txid, idx): out for txid, idx, out in utxos}    def resolver(txid, index):        return lookup.get((txid, index))    return resolver# --------------------------------------------------------------------------- ## Base58Check# --------------------------------------------------------------------------- #def test_base58check_round_trip():    for _ in range(20):        pkh = os.urandom(32).hex()        addr = address_from_pubkey_hash(pkh)        assert is_valid_address(addr)        assert pubkey_hash_from_address(addr) == pkhdef test_corrupted_address_fails_checksum():    pkh = os.urandom(32).hex()    addr = address_from_pubkey_hash(pkh)    # Flip one character to a different valid base58 char.    i = len(addr) // 2    original = addr[i]    replacement = "A" if original != "A" else "B"    corrupted = addr[:i] + replacement + addr[i + 1 :]    assert corrupted != addr    assert is_valid_address(corrupted) is False    with pytest.raises(ValueError):        pubkey_hash_from_address(corrupted)def test_leading_zero_round_trip():    # A pubkey hash whose leading bytes are zero must still round-trip.    pkh = "00" * 4 + os.urandom(28).hex()    addr = address_from_pubkey_hash(pkh)    assert pubkey_hash_from_address(addr) == pkhdef test_base58_leading_zero_bytes():    data = b"\x00\x00\x01\x02\x03"    assert base58_decode(base58_encode(data)) == data# --------------------------------------------------------------------------- ## Wallet key management# --------------------------------------------------------------------------- #def test_new_key_returns_valid_addresses_and_grows():    w = Wallet()    assert w.addresses == []    a1 = w.new_key()    assert is_valid_address(a1)    assert w.addresses == [a1]    a2 = w.new_key()    assert is_valid_address(a2)    assert w.addresses == [a1, a2]def test_new_key_addresses_differ():    w = Wallet()    a1 = w.new_key()    a2 = w.new_key()    assert a1 != a2def test_balance_sums_only_owned():    w = Wallet()    owned_addr = w.new_key()    owned_pkh = pubkey_hash_from_address(owned_addr)    # An address the wallet does NOT own.    other = Wallet()    other_pkh = pubkey_hash_from_address(other.new_key())    utxos = make_utxos_paying(        [(owned_pkh, 500), (other_pkh, 999), (owned_pkh, 250)]    )    assert w.balance(utxos) == 750# --------------------------------------------------------------------------- ## create_transaction# --------------------------------------------------------------------------- #def test_create_transaction_with_change():    w = Wallet()    src_addr = w.new_key()    src_pkh = pubkey_hash_from_address(src_addr)    recipient = Wallet()    to_addr = recipient.new_key()    utxos = make_utxos_paying([(src_pkh, 1000)])    before = set(w.addresses)    tx = w.create_transaction(utxos, to_addr, amount=600, fee=50)    resolver = make_resolver(utxos)    assert tx.verify(resolver) is True    assert tx.fee(resolver) == 50    # Recipient output.    to_pkh = pubkey_hash_from_address(to_addr)    assert tx.outputs[0].amount == 600    assert tx.outputs[0].pubkey_hash == to_pkh    # Change output: 1000 - 600 - 50 = 350 back to a newly generated address.    assert len(tx.outputs) == 2    assert tx.outputs[1].amount == 350    change_pkh = tx.outputs[1].pubkey_hash    assert w.owns(change_pkh)    new_addrs = set(w.addresses) - before    assert len(new_addrs) == 1    change_addr = new_addrs.pop()    assert pubkey_hash_from_address(change_addr) == change_pkh    # History recorded.    assert w.history[-1]["amount"] == 600    assert w.history[-1]["to"] == to_addr    assert w.history[-1]["txid"] == tx.txiddef test_create_transaction_exact_amount_no_change():    w = Wallet()    src_addr = w.new_key()    src_pkh = pubkey_hash_from_address(src_addr)    to_addr = Wallet().new_key()    utxos = make_utxos_paying([(src_pkh, 700)])    before = set(w.addresses)    tx = w.create_transaction(utxos, to_addr, amount=650, fee=50)    resolver = make_resolver(utxos)    assert tx.verify(resolver) is True    assert tx.fee(resolver) == 50    assert len(tx.outputs) == 1  # no change output    assert tx.outputs[0].amount == 650    # No new change key was generated (exact spend).    assert set(w.addresses) == beforedef test_create_transaction_insufficient_funds():    w = Wallet()    src_pkh = pubkey_hash_from_address(w.new_key())    to_addr = Wallet().new_key()    utxos = make_utxos_paying([(src_pkh, 100)])    with pytest.raises(ValueError):        w.create_transaction(utxos, to_addr, amount=200, fee=10)def test_create_transaction_multiple_inputs():    w = Wallet()    pkh1 = pubkey_hash_from_address(w.new_key())    pkh2 = pubkey_hash_from_address(w.new_key())    to_addr = Wallet().new_key()    utxos = make_utxos_paying([(pkh1, 400), (pkh2, 400)])    tx = w.create_transaction(utxos, to_addr, amount=700, fee=10)    resolver = make_resolver(utxos)    assert len(tx.inputs) == 2    assert tx.verify(resolver) is True    assert tx.fee(resolver) == 10