plm-lottery/tests/unit/test_broadcast.py

from datetime import datetime, timedelta, timezone
from types import SimpleNamespace

import pytest
from embit import script
from embit.bip32 import HDKey
from embit.transaction import Transaction, TransactionInput, TransactionOutput
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from app.config import settings
from app.db.base import Base
from app.db.models import PendingTransaction, User
from app.tx.broadcast import RbfError, bump_fee, should_bump
from app.wallet.plm_network import PLM_MAINNET
from app.wallet.psbt_builder import MAX_FEE_RATE_SAT_VB, Utxo, build_signed_transaction, estimate_vsize


def _key(seed_byte: int) -> HDKey:
    root = HDKey.from_seed(bytes([seed_byte]) * 32, version=PLM_MAINNET["xprv"])
    return root.derive("m/84h/746h/0h/0/0")


def test_should_bump_false_before_timeout():
    pending = PendingTransaction(
        kind="bet", current_txid="x", fee_rate_sat_vb=1, raw_tx_hex="00", status="pending",
        broadcast_at=datetime.now(timezone.utc), last_broadcast_at=datetime.now(timezone.utc),
    )
    assert should_bump(pending, datetime.now(timezone.utc), timeout_seconds=900) is False


def test_should_bump_true_after_timeout():
    pending = PendingTransaction(
        kind="bet", current_txid="x", fee_rate_sat_vb=1, raw_tx_hex="00", status="pending",
        broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        last_broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
    )
    assert should_bump(pending, datetime.now(timezone.utc), timeout_seconds=900) is True


def test_should_bump_false_when_not_pending():
    pending = PendingTransaction(
        kind="bet", current_txid="x", fee_rate_sat_vb=1, raw_tx_hex="00", status="confirmed",
        broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        last_broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
    )
    assert should_bump(pending, datetime.now(timezone.utc), timeout_seconds=900) is False


def test_should_bump_measures_from_last_broadcast_not_first(monkeypatch):
    """B-27 regression: a tx first broadcast long ago, but bumped recently, must not
    be due for another bump yet — should_bump has to look at last_broadcast_at, not
    the original broadcast_at, or every tick would try to re-bump it."""
    pending = PendingTransaction(
        kind="bet", current_txid="x", fee_rate_sat_vb=1, raw_tx_hex="00", status="pending",
        broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=10_000),
        last_broadcast_at=datetime.now(timezone.utc),
    )
    assert should_bump(pending, datetime.now(timezone.utc), timeout_seconds=900) is False


class FakeClient:
    """B-40: _prevout_amount now asks for the raw (non-verbose) transaction and
    reads its output value as an integer via embit, rather than a verbose reply's
    float "value" field — so this fake must hand back a real, parseable raw tx
    whose vout[0] carries the requested amount (every test here spends vout 0 of
    its fixture UTXO)."""

    def __init__(self, prevout_values: dict[str, int]):
        self._prevout_values = prevout_values
        self.broadcasted: list[str] = []

    async def get_transaction(self, txid: str, verbose: bool = False) -> str:
        assert verbose is False
        fake_prevout_tx = Transaction(
            vin=[TransactionInput(b"\x00" * 32, 0)],
            vout=[TransactionOutput(self._prevout_values[txid], script.Script(b"\x00\x14" + b"\x00" * 20))],
        )
        return fake_prevout_tx.serialize().hex()

    async def broadcast(self, raw_tx_hex: str) -> str:
        self.broadcasted.append(raw_tx_hex)
        return "network-txid"


@pytest.fixture
async def session_factory(tmp_path, monkeypatch):
    monkeypatch.setattr(settings, "master_key_path", str(tmp_path / "master.xprv.enc"))
    monkeypatch.setattr(
        settings,
        "xprv_encryption_key",
        __import__("cryptography.fernet", fromlist=["Fernet"]).Fernet.generate_key().decode(),
    )
    from app.wallet import hd

    hd._account_key = None
    hd.generate_master_key()

    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield async_sessionmaker(engine, expire_on_commit=False)
    await engine.dispose()
    hd._account_key = None


async def test_bump_fee_shrinks_change_and_rebroadcasts(session_factory):
    from app.wallet.hd import derive_user_address, derive_user_key

    signer = derive_user_key(0)
    my_address = derive_user_address(0)
    from_script = script.p2wpkh(signer.to_public())
    to_address = script.p2wpkh(_key(99).to_public()).address(network=PLM_MAINNET)

    utxo_amount = 150_000_000
    utxo_txid = "11" * 32
    built = build_signed_transaction(
        signing_key=signer,
        from_script=from_script,
        utxos=[Utxo(utxo_txid, 0, utxo_amount)],
        to_address=to_address,
        amount_sats=10_000_000,
        change_address=my_address,
        fee_rate_sat_vb=1,
    )

    async with session_factory() as session:
        user = User(username="alice", password_hash="x", derivation_index=0, address=my_address)
        session.add(user)
        await session.commit()
        pending = PendingTransaction(
            kind="bet",
            user_id=user.id,
            current_txid=built.txid,
            fee_rate_sat_vb=1,
            raw_tx_hex=built.raw_hex,
            status="pending",
            broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    client = FakeClient({utxo_txid: utxo_amount})

    new_txid = await bump_fee(session_factory, client, pending_id)

    assert client.broadcasted
    assert new_txid != built.txid

    new_tx = Transaction.parse(bytes.fromhex(client.broadcasted[0]))
    old_tx = Transaction.parse(bytes.fromhex(built.raw_hex))
    old_change = next(o.value for o in old_tx.vout if o.script_pubkey.address(network=PLM_MAINNET) == my_address)
    new_change = next(o.value for o in new_tx.vout if o.script_pubkey.address(network=PLM_MAINNET) == my_address)
    assert new_change < old_change  # fee bump came out of the change output

    async with session_factory() as session:
        row = await session.get(PendingTransaction, pending_id)
        assert row.current_txid == new_txid
        assert row.fee_rate_sat_vb == 2
        assert row.attempt_count == 2


async def test_bump_fee_leaves_broadcast_at_untouched(session_factory):
    """B-27 regression: bump_fee must only ever update last_broadcast_at. Before
    this, it overwrote broadcast_at on every bump — the same field
    tx/reconcile.py's abandon-after-N-hours grace period measures from — so a
    repeatedly-bumped-but-never-mined tx reset that clock forever and was never
    abandoned."""
    from app.wallet.hd import derive_user_address, derive_user_key

    signer = derive_user_key(0)
    my_address = derive_user_address(0)
    from_script = script.p2wpkh(signer.to_public())
    to_address = script.p2wpkh(_key(97).to_public()).address(network=PLM_MAINNET)

    utxo_amount = 150_000_000
    utxo_txid = "33" * 32
    built = build_signed_transaction(
        signing_key=signer,
        from_script=from_script,
        utxos=[Utxo(utxo_txid, 0, utxo_amount)],
        to_address=to_address,
        amount_sats=10_000_000,
        change_address=my_address,
        fee_rate_sat_vb=1,
    )

    original_broadcast_at = datetime.now(timezone.utc) - timedelta(days=1)
    async with session_factory() as session:
        user = User(username="carol", password_hash="x", derivation_index=0, address=my_address)
        session.add(user)
        await session.commit()
        pending = PendingTransaction(
            kind="bet",
            user_id=user.id,
            current_txid=built.txid,
            fee_rate_sat_vb=1,
            raw_tx_hex=built.raw_hex,
            status="pending",
            broadcast_at=original_broadcast_at,
            last_broadcast_at=original_broadcast_at,
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    client = FakeClient({utxo_txid: utxo_amount})
    before_bump = datetime.now(timezone.utc)

    await bump_fee(session_factory, client, pending_id)

    async with session_factory() as session:
        row = await session.get(PendingTransaction, pending_id)
        assert row.broadcast_at.replace(tzinfo=timezone.utc) == original_broadcast_at
        assert row.last_broadcast_at.replace(tzinfo=timezone.utc) >= before_bump


async def test_bump_fee_raises_when_no_change_output(session_factory):
    """The guard still matters after B-62 even though the builder no longer produces
    this shape: a single-output transaction broadcast before that change can still be
    sitting in `pending` across the deploy, and it must fail loudly rather than
    silently shrink the recipient's output. Hence a hand-built tx here — the point is
    exactly that build_signed_transaction won't make one any more."""
    from embit.transaction import Transaction, TransactionInput, TransactionOutput

    from app.wallet.hd import derive_user_address, derive_user_key
    from app.wallet.psbt_builder import RBF_SEQUENCE

    signer = derive_user_key(0)
    my_address = derive_user_address(0)
    to_address = script.p2wpkh(_key(98).to_public()).address(network=PLM_MAINNET)

    utxo_amount = 10_000_000  # entirely consumed by the single recipient output
    utxo_txid = "22" * 32
    legacy_tx = Transaction(
        vin=[TransactionInput(bytes.fromhex(utxo_txid), 0, sequence=RBF_SEQUENCE)],
        vout=[TransactionOutput(utxo_amount - 141, script.Script.from_address(to_address))],
    )
    built = SimpleNamespace(raw_hex=legacy_tx.serialize().hex(), txid=legacy_tx.txid().hex())

    async with session_factory() as session:
        user = User(username="bob", password_hash="x", derivation_index=0, address=my_address)
        session.add(user)
        await session.commit()
        pending = PendingTransaction(
            kind="bet",
            user_id=user.id,
            current_txid=built.txid,
            fee_rate_sat_vb=1,
            raw_tx_hex=built.raw_hex,
            status="pending",
            broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    client = FakeClient({utxo_txid: utxo_amount})

    with pytest.raises(RbfError):
        await bump_fee(session_factory, client, pending_id)


async def test_bump_fee_retargets_every_stored_txid(session_factory):
    """B-02/B-20: a bump changes the txid, and everything that recorded the old one
    has to follow — the participant's bet_txid (whose staleness used to wedge the
    round forever), the UTXO's spent_txid (which the reconciler matches on), and
    replaced_by_txid, which was never written at all."""
    from app.db.models import Round, RoundParticipant, UtxoEvent
    from app.wallet.hd import derive_user_address, derive_user_key

    signer = derive_user_key(0)
    my_address = derive_user_address(0)
    from_script = script.p2wpkh(signer.to_public())
    to_address = script.p2wpkh(_key(98).to_public()).address(network=PLM_MAINNET)

    utxo_amount = 150_000_000
    utxo_txid = "22" * 32
    built = build_signed_transaction(
        signing_key=signer,
        from_script=from_script,
        utxos=[Utxo(utxo_txid, 0, utxo_amount)],
        to_address=to_address,
        amount_sats=10_000_000,
        change_address=my_address,
        fee_rate_sat_vb=1,
    )

    async with session_factory() as session:
        user = User(username="bob", password_hash="x", derivation_index=0, address=my_address)
        session.add(user)
        session.add(Round(id=1, status="open"))
        await session.flush()
        session.add(
            UtxoEvent(
                user_id=user.id, txid=utxo_txid, vout=0, amount_sats=utxo_amount,
                confirmed_height=5, spent_txid=built.txid,
            )
        )
        session.add(
            RoundParticipant(
                round_id=1, user_id=user.id, bet_amount_sats=built.recipient_sats,
                bet_txid=built.txid, status="broadcast",
            )
        )
        pending = PendingTransaction(
            kind="bet", round_id=1, user_id=user.id, current_txid=built.txid, fee_rate_sat_vb=1,
            raw_tx_hex=built.raw_hex, status="pending",
            broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    new_txid = await bump_fee(session_factory, FakeClient({utxo_txid: utxo_amount}), pending_id)

    async with session_factory() as session:
        from sqlalchemy import select

        row = await session.get(PendingTransaction, pending_id)
        assert row.current_txid == new_txid
        assert row.replaced_by_txid == built.txid  # points backwards at what it replaced

        participant = (await session.scalars(select(RoundParticipant))).one()
        assert participant.bet_txid == new_txid

        utxo = (await session.scalars(select(UtxoEvent))).one()
        assert utxo.spent_txid == new_txid


# --- B-32: the bump delta must always meet BIP125's relay-mandated minimum, and
# escalation must stop at a ceiling instead of retrying forever. ------------------


async def test_bump_fee_meets_bip125_minimum_when_old_fee_already_exceeds_target(session_factory):
    """old_fee (as bump_fee computes it from the actual prevout amounts) can end
    up higher than vsize * target_fee_rate — e.g. because dust change was folded
    into the original fee (wallet/psbt_builder.py's DUST_LIMIT_SATS handling).
    The naive `target_fee - old_fee` goes negative in that case; the previous
    fallback was a flat 1-satoshi total bump, nowhere near BIP125 rule 4's
    required minimum, so the node rejected it every time and — since bump_fee
    raised before touching `pending` — the next tick retried identically every
    30 seconds, forever. Simulated here by reporting a prevout inflated beyond
    what was actually spent, which has the same effect on old_fee as dust
    absorption would."""
    from app.wallet.hd import derive_user_address, derive_user_key

    signer = derive_user_key(0)
    my_address = derive_user_address(0)
    from_script = script.p2wpkh(signer.to_public())
    to_address = script.p2wpkh(_key(96).to_public()).address(network=PLM_MAINNET)

    utxo_amount = 150_000_000
    utxo_txid = "55" * 32
    built = build_signed_transaction(
        signing_key=signer,
        from_script=from_script,
        utxos=[Utxo(utxo_txid, 0, utxo_amount)],
        to_address=to_address,
        amount_sats=10_000_000,
        change_address=my_address,
        fee_rate_sat_vb=1,
    )

    async with session_factory() as session:
        user = User(username="dave", password_hash="x", derivation_index=0, address=my_address)
        session.add(user)
        await session.commit()
        pending = PendingTransaction(
            kind="bet",
            user_id=user.id,
            current_txid=built.txid,
            fee_rate_sat_vb=1,
            raw_tx_hex=built.raw_hex,
            status="pending",
            broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    # Reports a prevout inflated well beyond what was actually spent — has the
    # same effect on old_fee as dust absorption would have: old_fee ends up far
    # above vsize * target_fee_rate (target_fee_rate = 2 here).
    inflated_excess = 50_000
    client = FakeClient({utxo_txid: utxo_amount + inflated_excess})

    new_txid = await bump_fee(session_factory, client, pending_id)

    assert client.broadcasted
    new_tx = Transaction.parse(bytes.fromhex(client.broadcasted[0]))
    old_tx = Transaction.parse(bytes.fromhex(built.raw_hex))
    old_change = next(o.value for o in old_tx.vout if o.script_pubkey.address(network=PLM_MAINNET) == my_address)
    new_change = next(o.value for o in new_tx.vout if o.script_pubkey.address(network=PLM_MAINNET) == my_address)

    vsize = estimate_vsize(len(old_tx.vin), len(old_tx.vout))
    min_valid_delta = vsize * 1  # BIP125 rule 4's floor at a 1 sat/vB incremental relay fee
    assert min_valid_delta > 1  # meaningfully more than the old flat "1 satoshi" fallback
    assert old_change - new_change == min_valid_delta

    async with session_factory() as session:
        row = await session.get(PendingTransaction, pending_id)
        old_fee_as_bump_fee_computed_it = built.fee_sats + inflated_excess
        expected_rate = (old_fee_as_bump_fee_computed_it + min_valid_delta) // vsize
        assert row.fee_rate_sat_vb == expected_rate
        assert row.fee_rate_sat_vb > 2  # the actual rate, not the naive (and too-low) target


async def test_bump_fee_refuses_once_at_the_max_fee_rate(session_factory):
    """Without a ceiling, a stuck transaction's fee rate climbed by 1 sat/vB every
    30 seconds forever, eating further and further into the user's change."""
    from app.wallet.hd import derive_user_address, derive_user_key

    signer = derive_user_key(0)
    my_address = derive_user_address(0)
    from_script = script.p2wpkh(signer.to_public())
    to_address = script.p2wpkh(_key(95).to_public()).address(network=PLM_MAINNET)

    utxo_amount = 150_000_000
    utxo_txid = "66" * 32
    built = build_signed_transaction(
        signing_key=signer,
        from_script=from_script,
        utxos=[Utxo(utxo_txid, 0, utxo_amount)],
        to_address=to_address,
        amount_sats=10_000_000,
        change_address=my_address,
        fee_rate_sat_vb=1,
    )

    async with session_factory() as session:
        user = User(username="erin", password_hash="x", derivation_index=0, address=my_address)
        session.add(user)
        await session.commit()
        pending = PendingTransaction(
            kind="bet",
            user_id=user.id,
            current_txid=built.txid,
            fee_rate_sat_vb=MAX_FEE_RATE_SAT_VB,
            raw_tx_hex=built.raw_hex,
            status="pending",
            broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    client = FakeClient({utxo_txid: utxo_amount})

    with pytest.raises(RbfError):
        await bump_fee(session_factory, client, pending_id)

    assert not client.broadcasted


# --- B-40: bump_fee must not hold a DB session open across its network calls,
# and a row that's no longer pending by the time it runs is a quiet no-op. -------


async def test_bump_fee_holds_no_session_open_during_network_calls(session_factory):
    """The get_transaction-per-input reads and the broadcast must happen with no
    DB session held open — the same shape used elsewhere for this reason (B-18,
    electrum/listener.py's refresh_user for B-31) — otherwise a session sits
    idle in the pool for the whole duration of what can be several slow network
    round-trips."""
    from app.wallet.hd import derive_user_address, derive_user_key

    signer = derive_user_key(0)
    my_address = derive_user_address(0)
    from_script = script.p2wpkh(signer.to_public())
    to_address = script.p2wpkh(_key(94).to_public()).address(network=PLM_MAINNET)

    utxo_amount = 150_000_000
    utxo_txid = "77" * 32
    built = build_signed_transaction(
        signing_key=signer,
        from_script=from_script,
        utxos=[Utxo(utxo_txid, 0, utxo_amount)],
        to_address=to_address,
        amount_sats=10_000_000,
        change_address=my_address,
        fee_rate_sat_vb=1,
    )

    async with session_factory() as session:
        user = User(username="frank", password_hash="x", derivation_index=0, address=my_address)
        session.add(user)
        await session.commit()
        pending = PendingTransaction(
            kind="bet",
            user_id=user.id,
            current_txid=built.txid,
            fee_rate_sat_vb=1,
            raw_tx_hex=built.raw_hex,
            status="pending",
            broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    open_count = {"n": 0}

    class _TrackedSession:
        def __init__(self, inner):
            self._inner = inner

        async def __aenter__(self):
            result = await self._inner.__aenter__()
            open_count["n"] += 1
            return result

        async def __aexit__(self, *exc):
            open_count["n"] -= 1
            return await self._inner.__aexit__(*exc)

    def tracking_session_factory():
        return _TrackedSession(session_factory())

    class TrackingClient(FakeClient):
        async def get_transaction(self, txid, verbose=False):
            assert open_count["n"] == 0, "a session was held open during a network call"
            return await super().get_transaction(txid, verbose)

        async def broadcast(self, raw_tx_hex):
            assert open_count["n"] == 0, "a session was held open during the broadcast"
            return await super().broadcast(raw_tx_hex)

    client = TrackingClient({utxo_txid: utxo_amount})
    await bump_fee(tracking_session_factory, client, pending_id)

    assert client.broadcasted
    assert open_count["n"] == 0  # nothing left open afterwards either


async def test_bump_fee_is_a_noop_when_no_longer_pending(session_factory):
    """A row can legitimately confirm (or otherwise leave "pending") between
    being read as due and RbfBumper actually attempting the bump — a normal
    race, not an error. Must return quietly rather than raising or touching
    the network."""
    async with session_factory() as session:
        user = User(username="grace", password_hash="x", derivation_index=0, address="plm1qxxx")
        session.add(user)
        await session.commit()
        pending = PendingTransaction(
            kind="bet",
            user_id=user.id,
            current_txid="already-confirmed-txid",
            fee_rate_sat_vb=1,
            raw_tx_hex="00",
            status="confirmed",
            broadcast_at=datetime.now(timezone.utc) - timedelta(seconds=1000),
        )
        session.add(pending)
        await session.commit()
        pending_id = pending.id

    client = FakeClient({})

    result = await bump_fee(session_factory, client, pending_id)

    assert result is None
    assert not client.broadcasted


async def test_bump_fee_is_a_noop_when_the_row_is_gone(session_factory):
    client = FakeClient({})
    result = await bump_fee(session_factory, client, 999_999)
    assert result is None
    assert not client.broadcasted