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.config import settings
from app.db.base import Base
from app.db.models import PendingTransaction, User, UtxoEvent, Withdrawal
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)


@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_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