plm-lottery/tests/unit/test_round_pool_address.py

"""B-80: every round has its own pool address (RND-35, DAT-60, DAT-11, REG-14, SEC-08,
BET-12, RND-22, WDR-14, API-23, API-36).

All bets of all rounds used to pay one address `m/84'/746'/0'/1/0`: its history grew
without bound, other rounds' money and stray coins mixed with a round's own, and the
payout picked from everything it held. Now round N pays `m/84'/746'/0'/2/<N>`, stored
on the round; the old shared address survives as the *reserve* address and as the pool
of the rounds opened before the change.
"""

import os
import sqlite3
import subprocess
import sys
from pathlib import Path

import pytest
from embit import script
from embit.ec import PublicKey
from embit.transaction import Transaction
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 AuditLog, PendingTransaction, Round, RoundConfig, User, UtxoEvent, Withdrawal
from app.electrum.scripthash import address_to_scripthash
from app.rounds.scheduler import RoundScheduler
from app.rounds.service import open_new_round_if_needed
from app.wallet import hd
from app.wallet.hd import derive_pool_address, derive_pool_key, derive_round_pool_key, derive_user_address
from app.wallet.plm_network import PLM_MAINNET
from app.withdrawals.service import WithdrawalError, request_withdrawal
from tests.unit.test_scheduler import (  # noqa: F401  (fixture + helpers shared with the payout tests)
    _FEE_ADDRESS,
    _POOL_AMOUNT_SATS,
    FakePayoutListener,
    _seed_paying_out_round,
    payout_session_factory,
)

_ROOT = Path(__file__).resolve().parents[2]
_EXTERNAL = "plm1qqph9qup2mp7w7g5nlsdhdc9m2pp44ampzw0ctx"


def _address_of(key) -> str:
    return script.p2wpkh(key.to_public()).address(network=PLM_MAINNET)


# --- derivation ---------------------------------------------------------------


def test_each_round_has_its_own_branch_2_address_and_the_reserve_is_branch_1():
    """SEC-08 / REG-14: round N is m/84'/746'/0'/2/N, no round id is the reserve 1/0."""
    account = hd._load_account_key()
    assert _address_of(derive_pool_key(5)) == _address_of(account.derive("2/5"))
    assert derive_pool_address(5) == _address_of(account.derive("2/5"))
    assert derive_pool_address() == _address_of(account.derive("1/0"))  # the reserve
    assert len({derive_pool_address(1), derive_pool_address(2), derive_pool_address()}) == 3
    with pytest.raises(ValueError):
        derive_pool_key(0)


def test_the_signing_key_must_control_the_stored_address():
    """DAT-60 / RND-22: a round's own address is signed with its own key, a round opened
    before B-80 (reserve address stored) with the reserve key, anything else is refused."""
    assert _address_of(derive_round_pool_key(3, derive_pool_address(3))) == derive_pool_address(3)
    legacy = derive_round_pool_key(3, derive_pool_address())  # id 3, but the shared address
    assert _address_of(legacy) == derive_pool_address()
    with pytest.raises(ValueError):
        derive_round_pool_key(3, derive_pool_address(4))  # another round's address
    with pytest.raises(ValueError):
        derive_round_pool_key(3, "")


# --- opening a round ----------------------------------------------------------


@pytest.fixture
async def session_factory():
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    async with async_sessionmaker(engine, expire_on_commit=False)() as session:
        session.add(RoundConfig(fee_address=_FEE_ADDRESS, round_cooldown_seconds=0))
        await session.commit()
    yield async_sessionmaker(engine, expire_on_commit=False)
    await engine.dispose()


async def test_opening_a_round_stores_its_own_pool_address(session_factory):
    """RND-35 / DAT-60: written in the transaction that inserts the round, from its id."""
    async with session_factory() as session:
        first = await open_new_round_if_needed(session)
        await session.commit()
        first_id, first_address = first.id, first.pool_address
        first.status = "closed"
        first.closed_at = first.opened_at
        await session.commit()

    async with session_factory() as session:
        second = await open_new_round_if_needed(session)
        await session.commit()
        second_id, second_address = second.id, second.pool_address

    assert first_id != second_id
    assert first_address == derive_pool_address(first_id)
    assert second_address == derive_pool_address(second_id)
    assert first_address != second_address
    assert "" not in (first_address, second_address)


# --- bets ---------------------------------------------------------------------


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

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


async def _funded_user(session_factory, index: int, sats: int) -> int:
    async with session_factory() as session:
        user = User(
            username=f"user{index}", password_hash="x", derivation_index=index, address=derive_user_address(index)
        )
        session.add(user)
        await session.commit()
        session.add(UtxoEvent(user_id=user.id, txid=f"{index + 1:02x}" * 32, vout=0, amount_sats=sats, confirmed_height=1))
        await session.commit()
        return user.id


async def test_a_bet_pays_the_pool_address_of_its_round(session_factory):
    """BET-12 / B-80: the bet's recipient output is round.pool_address, not the reserve."""
    user_id = await _funded_user(session_factory, 0, 2_000_000_000)
    client = _RecordingClient()
    async with session_factory() as session:
        user = await session.get(User, user_id)
        await place_bet(session, client, user)
        round_ = (await session.scalars(select(Round))).one()
        pool_address = round_.pool_address

    assert pool_address == derive_pool_address(round_.id)
    outputs = {o.script_pubkey.address(network=PLM_MAINNET) for o in Transaction.parse(bytes.fromhex(client.broadcasted[0])).vout}
    assert pool_address in outputs
    assert derive_pool_address() not in outputs  # the reserve address gets no bets any more


# --- payout -------------------------------------------------------------------


class _RecordingPayoutClient:
    def __init__(self, entries):
        self._entries = entries
        self.scripthashes: list[str] = []
        self.broadcasted: list[str] = []

    async def listunspent(self, scripthash):
        self.scripthashes.append(scripthash)
        return self._entries

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


def _signer_address(raw_tx_hex: str) -> str:
    witness = Transaction.parse(bytes.fromhex(raw_tx_hex)).vin[0].witness.items
    return script.p2wpkh(PublicKey.parse(witness[1])).address(network=PLM_MAINNET)


async def _payout(payout_session_factory, pool_address: str, round_id: int = 1):
    await _seed_paying_out_round(payout_session_factory, round_id)
    async with payout_session_factory() as session:
        round_ = await session.get(Round, round_id)
        round_.pool_address = pool_address
        await session.commit()
    client = _RecordingPayoutClient(
        [{"tx_hash": "44" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS}]
    )
    listener = FakePayoutListener(client)
    await RoundScheduler(payout_session_factory, listener)._trigger_payout(round_id)
    # B-82: a round with its own address asks the quorum (the listener), not the active
    # connection, which is only asked for a round paid from the shared reserve address.
    client.scripthashes = client.scripthashes + listener.verified_scripthashes
    return client


async def test_the_payout_reads_and_signs_with_the_round_own_pool_address(payout_session_factory):
    """RND-22 / DAT-11 / B-80: only the round's address is asked about, and its key signs."""
    client = await _payout(payout_session_factory, derive_pool_address(1))
    assert client.scripthashes == [address_to_scripthash(derive_pool_address(1))]
    assert _signer_address(client.broadcasted[0]) == derive_pool_address(1)


async def test_a_round_opened_before_b80_is_still_paid_from_the_reserve_address(payout_session_factory):
    """DAT-60 / RND-22: its stored pool_address is the old shared one, signed with 1/0."""
    client = await _payout(payout_session_factory, derive_pool_address(), round_id=6)
    assert client.scripthashes == [address_to_scripthash(derive_pool_address())]
    assert _signer_address(client.broadcasted[0]) == derive_pool_address()


@pytest.mark.parametrize("stored", ["", "plm1q5x25wd6q463mfhckjraaedgjg0lyu73qfcj43n"])
async def test_a_payout_never_signs_for_an_address_its_keys_do_not_control(payout_session_factory, stored):
    """RND-22 / B-80: an empty or foreign stored address builds nothing and sends nothing."""
    client = await _payout(payout_session_factory, stored)
    assert client.broadcasted == []
    assert client.scripthashes == []
    async with payout_session_factory() as session:
        reasons = (await session.scalars(select(AuditLog.payload_json).where(AuditLog.event_type == "payout_failed"))).all()
        assert len(reasons) == 1 and "payout build failed" in reasons[0]


async def test_an_rbf_bump_of_a_payout_uses_the_round_own_key(payout_session_factory):
    """TX-11 / B-80: tx/broadcast.py:_signing_context follows the round's pool address."""
    from app.tx.broadcast import RbfError, _signing_context

    await _seed_paying_out_round(payout_session_factory, 4)
    async with payout_session_factory() as session:
        round_ = await session.get(Round, 4)
        round_.pool_address = derive_pool_address(4)
        await session.commit()
        _key, _script, address = await _signing_context(session, "payout", None, 4)
        assert address == derive_pool_address(4)

        round_.pool_address = ""
        with pytest.raises(RbfError):
            await _signing_context(session, "payout", None, 4)


# --- withdrawals --------------------------------------------------------------


@pytest.mark.parametrize("which", ["round_pool", "reserve", "fee"])
async def test_a_withdrawal_to_an_address_of_the_platform_is_refused(session_factory, which):
    """WDR-14 / SEC-43 / B-80: refused with withdrawal_to_platform_address before any coin
    is selected or marked spent."""
    user_id = await _funded_user(session_factory, 3, 3_000_000_000)
    async with session_factory() as session:
        round_ = await open_new_round_if_needed(session)
        await session.commit()
        destination = {
            "round_pool": round_.pool_address,
            "reserve": derive_pool_address(),
            "fee": _FEE_ADDRESS,
        }[which]

    client = _RecordingClient()
    async with session_factory() as session:
        user = await session.get(User, user_id)
        with pytest.raises(WithdrawalError) as excinfo:
            await request_withdrawal(session, client, user, destination, 1_000_000_000)
        assert excinfo.value.code == "withdrawal_to_platform_address"
        # Case does not get around it: bech32 is case-insensitive.
        with pytest.raises(WithdrawalError) as excinfo:
            await request_withdrawal(session, client, user, destination.upper(), 1_000_000_000)
        assert excinfo.value.code in ("withdrawal_to_platform_address", "invalid_address")

    assert client.broadcasted == []
    async with session_factory() as session:
        assert (await session.scalars(select(UtxoEvent.spent_txid))).all() == [None]
        assert (await session.scalars(select(Withdrawal))).all() == []


async def test_the_old_pool_addresses_of_closed_rounds_stay_refused(session_factory):
    """WDR-14: a *closed* round's address still belongs to the platform."""
    user_id = await _funded_user(session_factory, 3, 3_000_000_000)
    async with session_factory() as session:
        session.add(Round(id=40, status="closed", pool_address=derive_pool_address(40)))
        await session.commit()
    async with session_factory() as session:
        user = await session.get(User, user_id)
        with pytest.raises(WithdrawalError) as excinfo:
            await request_withdrawal(session, _RecordingClient(), user, derive_pool_address(40), 1_000_000_000)
        assert excinfo.value.code == "withdrawal_to_platform_address"


def test_the_platform_address_error_is_translated_in_every_language():
    """The `error.withdrawal_to_platform_address` code needs all 7 languages (CLAUDE.md)."""
    source = (_ROOT / "app" / "static" / "i18n.js").read_text(encoding="utf-8")
    assert source.count("'error.withdrawal_to_platform_address'") == 7


# --- migration ----------------------------------------------------------------


def test_the_migration_backfills_existing_rounds_with_the_reserve_address(tmp_path):
    """DAT-60 / B-80: rounds that exist at upgrade time all paid the shared address, so they
    get the reserve address as pool_address; the expression indexes survive (B-97)."""
    key_path = tmp_path / "master.xprv.enc"
    db_path = tmp_path / "old.db"
    # The autouse fixture already put a throwaway key in `settings`: write it where the
    # subprocess will read it and hand it the same encryption key.
    key_path.write_bytes(Path(settings.master_key_path).read_bytes())
    env = {
        **os.environ,
        "DATABASE_URL": f"sqlite+aiosqlite:///{db_path}",
        "MASTER_KEY_PATH": str(key_path),
        "XPRV_ENCRYPTION_KEY": settings.xprv_encryption_key,
    }

    def alembic(*args):
        result = subprocess.run(
            [sys.executable, "-m", "alembic", *args], cwd=_ROOT, env=env, capture_output=True, text=True, timeout=120
        )
        assert result.returncode == 0, result.stderr[-3000:]

    alembic("upgrade", "283844a44b4a")  # the revision just before pool_address
    with sqlite3.connect(db_path) as conn:
        conn.execute("INSERT INTO rounds (status, opened_at) VALUES ('closed', '2026-01-01 00:00:00')")
        conn.execute("INSERT INTO rounds (status, opened_at) VALUES ('open', '2026-01-02 00:00:00')")
    alembic("upgrade", "head")

    with sqlite3.connect(db_path) as conn:
        addresses = [row[0] for row in conn.execute("SELECT pool_address FROM rounds ORDER BY id")]
        indexes = {row[0] for row in conn.execute("SELECT name FROM sqlite_master WHERE type = 'index'")}
    assert addresses == [derive_pool_address()] * 2
    assert {"ix_rounds_single_active", "ix_users_username_lower"} <= indexes