plm-lottery/tests/unit/test_balance.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
from app.wallet.balance import compute_pending_balance, recompute_balance
from app.wallet.hd import derive_user_address


class FakeElectrumClient:
    async def broadcast(self, raw_tx_hex: str) -> str:
        return "fake-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 _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 recompute_balance(session, user.id)
        await session.commit()
        return user.id


async def test_pending_balance_includes_unconfirmed_change(session_factory):
    """A bet spends a whole (much larger) UTXO and the change hasn't confirmed
    yet, so cached_balance_sats alone understates the user's real balance by
    the entire unconfirmed change amount — compute_pending_balance should add
    it back."""
    user_id = await _make_funded_user(session_factory, 0, 1_500_000_000)
    client = FakeElectrumClient()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        await place_bet(session, 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

        pending_balance, has_pending = await compute_pending_balance(session, user)

        assert has_pending is True
        # confirmed (0) + unconfirmed change should be just under the original
        # funding amount (minus the bet amount and the network fee)
        assert 0 < pending_balance < 1_500_000_000


async def test_pending_balance_matches_confirmed_when_nothing_in_flight(session_factory):
    user_id = await _make_funded_user(session_factory, 1, 2_000_000_000)

    async with session_factory() as session:
        user = await session.get(User, user_id)
        pending_balance, has_pending = await compute_pending_balance(session, user)

    assert has_pending is False
    assert pending_balance == 2_000_000_000


async def test_pending_balance_does_not_double_count_change_already_credited(session_factory):
    """The Electrum listener (event-driven) and the confirmation poller (10s
    cadence) independently react to the same change output confirming. When the
    listener wins that race — the common case — the change is already a
    UtxoEvent inside cached_balance_sats while the PendingTransaction row is
    still "pending". compute_pending_balance must not add the change a second
    time in that window."""
    user_id = await _make_funded_user(session_factory, 4, 1_500_000_000)
    client = FakeElectrumClient()

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

    async with session_factory() as session:
        pending = (await session.scalars(select(PendingTransaction))).one()
        from embit.transaction import Transaction

        from app.wallet.plm_network import PLM_MAINNET

        tx = Transaction.parse(bytes.fromhex(pending.raw_tx_hex))
        change_vout, change_out = next(
            (i, out) for i, out in enumerate(tx.vout) if out.script_pubkey.address(network=PLM_MAINNET) == user.address
        )

        user = await session.get(User, user_id)
        # Simulate the listener having already credited the change output as
        # confirmed, before the poller has flipped `pending.status`.
        session.add(
            UtxoEvent(
                user_id=user_id,
                txid=pending.current_txid,
                vout=change_vout,
                amount_sats=change_out.value,
                confirmed_height=101,
            )
        )
        await recompute_balance(session, user_id)
        await session.commit()

    async with session_factory() as session:
        user = await session.get(User, user_id)
        pending_balance, has_pending = await compute_pending_balance(session, user)

        assert has_pending is True  # the PendingTransaction row is still "pending"
        assert pending_balance == user.cached_balance_sats  # already-credited change isn't added again


async def test_pending_balance_ignores_other_users_pending_transactions(session_factory):
    user_id = await _make_funded_user(session_factory, 2, 2_000_000_000)
    other_user_id = await _make_funded_user(session_factory, 3, 1_500_000_000)
    client = FakeElectrumClient()

    async with session_factory() as session:
        other_user = await session.get(User, other_user_id)
        await place_bet(session, client, other_user)

    async with session_factory() as session:
        pending_rows = (await session.scalars(select(PendingTransaction))).all()
        assert len(pending_rows) == 1  # sanity: only the other user has anything in flight

        user = await session.get(User, user_id)
        pending_balance, has_pending = await compute_pending_balance(session, user)

    assert has_pending is False
    assert pending_balance == 2_000_000_000