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
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)