plm-lottery/tests/unit/test_withdrawals.py

import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from app.bets.service import place_bet
from app.config import settings
from app.db.base import Base
from app.db.models import PendingTransaction, User, UtxoEvent, Withdrawal
from app.rounds.events import broadcaster
from app.wallet.hd import derive_user_address
from app.withdrawals.service import WithdrawalError, request_withdrawal


class FakeElectrumClient:
    def __init__(self):
        self.broadcasted: list[str] = []

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


EXTERNAL_ADDRESS = "plm1qqph9qup2mp7w7g5nlsdhdc9m2pp44ampzw0ctx"
BET_AMOUNT_SATS = 1_000_000_000  # matches RoundConfig.bet_amount_sats' column default; also the withdrawal minimum


@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 _make_funded_user(session_factory, index: int, funded_sats: int) -> int:
    async with session_factory() as session:
        address = derive_user_address(index)
        user = User(username=f"user{index}", password_hash="x", derivation_index=index, address=address)
        session.add(user)
        await session.commit()
        session.add(
            UtxoEvent(user_id=user.id, txid=f"{index:02x}" * 32, vout=0, amount_sats=funded_sats, confirmed_height=100)
        )
        await session.commit()
        return user.id


async def test_withdrawal_broadcasts_and_updates_balance(session_factory):
    user_id = await _make_funded_user(session_factory, 0, 2_000_000_000)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        withdrawal = await request_withdrawal(session, client, user, EXTERNAL_ADDRESS, BET_AMOUNT_SATS)

    assert client.broadcasted
    assert withdrawal.status == "broadcast"
    assert withdrawal.amount_sent_sats < BET_AMOUNT_SATS  # fee deducted from the amount

    async with session_factory() as session:
        user = await session.get(User, user_id)
        # The spent UTXO is gone immediately; the change output isn't credited
        # until it's independently observed as confirmed on-chain (same as bets) —
        # so the cached balance is transiently 0 until then, not the pre-fee delta.
        assert user.cached_balance_sats == 0
        pending = (await session.scalars(select(PendingTransaction))).one()
        assert pending.kind == "withdrawal"
        assert pending.withdrawal_id == withdrawal.id


async def test_withdrawal_rejects_amount_below_minimum(session_factory):
    user_id = await _make_funded_user(session_factory, 1, 2_000_000_000)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        with pytest.raises(WithdrawalError, match="minimum"):
            await request_withdrawal(session, client, user, EXTERNAL_ADDRESS, BET_AMOUNT_SATS - 1)


async def test_withdrawal_rejects_insufficient_balance(session_factory):
    user_id = await _make_funded_user(session_factory, 2, 1_000_000)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        with pytest.raises(WithdrawalError, match="insufficient balance"):
            await request_withdrawal(session, client, user, EXTERNAL_ADDRESS, BET_AMOUNT_SATS)


async def test_withdrawal_distinguishes_pending_from_truly_insufficient_balance(session_factory):
    """B-37: right after a bet, cached_balance_sats is ~0 because the whole funding
    UTXO was spent as input and the change hasn't confirmed yet — but the UI shows
    the pending-inclusive balance (compute_pending_balance), which does cover a
    withdrawal of this size. The error must say "not confirmed yet", not flatly
    "insufficient balance", or it contradicts what the user is looking at."""
    user_id = await _make_funded_user(session_factory, 4, 3_000_000_000)
    bet_client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        await place_bet(session, bet_client, user)

    async with session_factory() as session:
        user = await session.get(User, user_id)
        assert user.cached_balance_sats == 0  # the whole funding UTXO was spent as input

        withdraw_client = FakeElectrumClient()
        with pytest.raises(WithdrawalError) as exc_info:
            # Above the withdrawal minimum (BET_AMOUNT_SATS) and covered by the
            # unconfirmed change (~1_999_800_000 sats), but not by the (zero)
            # confirmed balance.
            await request_withdrawal(session, withdraw_client, user, EXTERNAL_ADDRESS, 1_500_000_000)

    assert exc_info.value.code == "balance_pending_confirmation"
    assert exc_info.value.params["pending_sats"] > 0
    assert not withdraw_client.broadcasted


@pytest.mark.parametrize(
    "address",
    [
        "bc1qw508d6qejxtdg4y5r3zarvary0c5xw7kv8f3t4",  # valid bech32, wrong chain
        "plm1qbogus",  # right HRP, broken checksum
        "not-an-address",
    ],
)
async def test_withdrawal_rejects_non_plm_address(session_factory, address):
    """The bc1 case is the one that matters: embit parses it into a perfectly
    valid witness program, so without the HRP check the withdrawal would build,
    sign and broadcast on PLM, sending the funds somewhere nobody holds a key
    for. It has to fail before a single UTXO is touched."""
    user_id = await _make_funded_user(session_factory, 3, 2_000_000_000)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        with pytest.raises(WithdrawalError) as exc_info:
            await request_withdrawal(session, client, user, address, BET_AMOUNT_SATS)

    assert exc_info.value.code == "invalid_address"
    assert not client.broadcasted


async def test_withdrawal_to_own_address_is_rejected(session_factory):
    """B-17: allowed before, and it broke two things that assume the recipient and
    the change are distinguishable by address — the RBF bump would shrink the
    recipient output, and compute_pending_balance counted the amount twice."""
    user_id = await _make_funded_user(session_factory, 8, 3_000_000_000)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        with pytest.raises(WithdrawalError, match="own deposit address"):
            await request_withdrawal(session, client, user, user.address, 1_000_000_000)

    assert not client.broadcasted
    async with session_factory() as session:
        assert (await session.scalars(select(UtxoEvent).where(UtxoEvent.user_id == user_id))).one().spent_txid is None


async def test_failed_broadcast_publishes_an_sse_update(session_factory):  # B-49
    """The released UTXOs are spendable again and the balance changed back, so the
    rollback must nudge the dashboard to refetch instead of leaving it stale until
    its next poll."""
    user_id = await _make_funded_user(session_factory, 10, 3_000_000_000)

    class RejectingClient:
        async def broadcast(self, raw_tx_hex: str) -> str:
            raise RuntimeError("min relay fee not met")

    queue = broadcaster.subscribe()
    try:
        while not queue.empty():
            queue.get_nowait()

        async with session_factory() as session:
            user = await session.get(User, user_id)
            with pytest.raises(WithdrawalError, match="refused"):
                await request_withdrawal(session, RejectingClient(), user, derive_user_address(98), 1_000_000_000)

        assert not queue.empty()
    finally:
        broadcaster.unsubscribe(queue)


async def test_failed_broadcast_marks_the_withdrawal_failed_and_frees_the_coins(session_factory):
    """B-07/B-08: the Withdrawal row is kept (unlike a bet) so the user can see the
    instruction didn't go through, but the coins must come back."""
    user_id = await _make_funded_user(session_factory, 9, 3_000_000_000)

    class RejectingClient:
        async def broadcast(self, raw_tx_hex: str) -> str:
            raise RuntimeError("min relay fee not met")

    async with session_factory() as session:
        user = await session.get(User, user_id)
        external = derive_user_address(99)
        with pytest.raises(WithdrawalError, match="refused"):
            await request_withdrawal(session, RejectingClient(), user, external, 1_000_000_000)

    async with session_factory() as session:
        withdrawal = (await session.scalars(select(Withdrawal))).one()
        assert withdrawal.status == "failed"
        assert withdrawal.txid is None
        assert (await session.scalars(select(UtxoEvent).where(UtxoEvent.user_id == user_id))).one().spent_txid is None
        user = await session.get(User, user_id)
        assert user.cached_balance_sats == 3_000_000_000


# --- B-62: "withdraw everything" must not build an unbumpable transaction ---------


async def test_full_balance_withdrawal_keeps_a_bumpable_change_output(session_factory):
    """The UI's max-amount checkbox sends the whole confirmed balance, so change came
    out at 0, the change output was dropped, and the tx had a single output —
    bump_fee then had nothing to shrink and raised RbfError every 30s until the
    reconciler abandoned the row hours later. Adding inputs is no answer here: the tx
    already spends every UTXO the user has. So a dust limit stays behind instead."""
    from embit.transaction import Transaction

    from app.wallet.psbt_builder import DUST_LIMIT_SATS

    balance = 2_000_000_000
    user_id = await _make_funded_user(session_factory, 40, balance)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        withdrawal = await request_withdrawal(session, client, user, EXTERNAL_ADDRESS, balance)

    tx = Transaction.parse(bytes.fromhex(client.broadcasted[0]))
    assert len(tx.vout) == 2  # recipient + change: bumpable
    assert all(o.value >= DUST_LIMIT_SATS for o in tx.vout)

    # The user asked for everything and is told what actually went out — the row
    # already distinguishes the two, since the fee comes out of the amount anyway.
    assert withdrawal.amount_requested_sats == balance
    fee = balance - sum(o.value for o in tx.vout)
    change = min(o.value for o in tx.vout)
    assert change == DUST_LIMIT_SATS
    assert withdrawal.amount_sent_sats == balance - DUST_LIMIT_SATS - fee


async def test_a_bet_from_a_balance_equal_to_the_bet_is_refused(session_factory):
    """The same shape on the PLAY side, where reducing the amount isn't an option —
    the bet is a fixed price. "A user's balance must never exactly equal the bet" is
    a documented invariant of the PLAY phase; this is where it's enforced, with an
    error that says how much more is needed rather than a bare "insufficient"."""
    from app.bets.service import BetError
    from app.wallet.psbt_builder import DUST_LIMIT_SATS

    user_id = await _make_funded_user(session_factory, 41, BET_AMOUNT_SATS)  # exactly the bet
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        with pytest.raises(BetError) as excinfo:
            await place_bet(session, client, user)

    assert excinfo.value.code == "balance_leaves_no_change"
    assert excinfo.value.params == {"required_extra_sats": DUST_LIMIT_SATS}
    assert not client.broadcasted

    async with session_factory() as session:
        utxo = (await session.scalars(select(UtxoEvent).where(UtxoEvent.user_id == user_id))).one()
        assert utxo.spent_txid is None  # refused before anything moved


async def test_a_bet_with_a_dust_limit_of_headroom_is_accepted(session_factory):
    from app.wallet.psbt_builder import DUST_LIMIT_SATS

    user_id = await _make_funded_user(session_factory, 42, BET_AMOUNT_SATS + DUST_LIMIT_SATS)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        participant = await place_bet(session, client, user)

    assert participant.status == "broadcast"