plm-lottery/tests/unit/test_scheduler.py

from datetime import datetime, timedelta, timezone

import pytest
from embit.transaction import Transaction
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 AuditLog, PendingTransaction, Round, RoundConfig, User
from app.rounds.scheduler import RoundScheduler, _reserved_payout_outpoints
from app.wallet.psbt_builder import MAX_TX_INPUTS


class FakeListener:
    client = object()  # truthy sentinel; _tick only checks "is not None"
    tip_height = 100
    tip_header_hex = "00"


@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)
    yield async_sessionmaker(engine, expire_on_commit=False)
    await engine.dispose()


async def test_tick_survives_sqlite_naive_datetime_roundtrip(session_factory, monkeypatch):
    """Regression test: SQLite drops tzinfo on round-trip, so opened_at comes back
    naive even though it was written as an aware UTC datetime. A prior bug compared
    it directly against datetime.now(timezone.utc) and crashed with
    "can't compare offset-naive and offset-aware datetimes" on every tick once a
    round existed — this must not happen."""
    async with session_factory() as session:
        session.add(RoundConfig(fee_address="", round_duration_seconds=3600))  # not due yet
        session.add(Round(status="open", opened_at=datetime.now(timezone.utc)))
        await session.commit()

    scheduler = RoundScheduler(session_factory, FakeListener())
    await scheduler._tick()  # must not raise


async def test_tick_closes_round_with_no_participants_once_due(session_factory, monkeypatch):
    past = datetime.now(timezone.utc) - timedelta(seconds=10)
    async with session_factory() as session:
        session.add(RoundConfig(fee_address="", round_duration_seconds=1))
        session.add(Round(status="open", opened_at=past))
        await session.commit()

    scheduler = RoundScheduler(session_factory, FakeListener())
    await scheduler._tick()

    async with session_factory() as session:
        round_ = (await session.scalars(select(Round))).one()
        assert round_.status == "closed"


# --- B-25: the payout must be persisted before it is broadcast, like bets/withdrawals ---

# A real, reusable PLM bech32 address so build_payout_transaction's
# script.Script.from_address(...) succeeds — this is not a value the scheduler
# validates itself (that's the admin panel's job for fee_address), it just needs to
# actually decode.
_WINNER_ADDRESS = "plm1q0xcqpzrky6eff2g52qdye53xkk9jxkvraxkkwd"
_FEE_ADDRESS = "plm1q5x25wd6q463mfhckjraaedgjg0lyu73qfcj43n"

_POOL_AMOUNT_SATS = 10_000_000_000  # 100 PLM


class FakePayoutClient:
    def __init__(self, entries, *, fail_broadcast=False):
        self._entries = entries
        self._fail_broadcast = fail_broadcast
        self.broadcasted: list[str] = []

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

    async def broadcast(self, raw_tx_hex):
        if self._fail_broadcast:
            raise RuntimeError("node rejected the transaction")
        self.broadcasted.append(raw_tx_hex)
        return "network-txid"


class FakePayoutListener:
    def __init__(self, client):
        self.client = client


@pytest.fixture
async def payout_session_factory(tmp_path, monkeypatch):
    """Same master-key bootstrap as test_broadcast.py's fixture: _trigger_payout
    needs a real pool key to sign with."""
    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 _seed_paying_out_round(session_factory, round_id: int = 1) -> int:
    async with session_factory() as session:
        winner = User(username="winner", password_hash="x", derivation_index=0, address=_WINNER_ADDRESS)
        session.add(winner)
        await session.flush()
        session.add(RoundConfig(fee_address=_FEE_ADDRESS, fee_rate_sat_vb=1))
        session.add(
            Round(
                id=round_id,
                status="paying_out",
                pool_amount_sats=_POOL_AMOUNT_SATS,
                winner_user_id=winner.id,
            )
        )
        await session.commit()
        return winner.id


async def test_trigger_payout_persists_before_broadcasting(payout_session_factory):
    """The happy path: payout_txid and a PendingTransaction must exist once the
    broadcast succeeds, promoted from "building" to "pending" — the two-phase write
    that used to be missing entirely (B-25)."""
    await _seed_paying_out_round(payout_session_factory)
    entries = [{"tx_hash": "33" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS + 100_000}]
    client = FakePayoutClient(entries)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client))

    await scheduler._trigger_payout(1)

    assert client.broadcasted
    async with payout_session_factory() as session:
        round_ = await session.get(Round, 1)
        assert round_.payout_txid is not None
        assert round_.winner_amount_sats and round_.fee_amount_sats

        pending = (await session.scalars(select(PendingTransaction))).one()
        assert pending.kind == "payout"
        assert pending.status == "pending"
        assert pending.current_txid == round_.payout_txid

        events = [e.event_type for e in (await session.scalars(select(AuditLog))).all()]
        assert "payout_sent" in events


async def test_trigger_payout_pays_a_round_with_more_participants_than_max_tx_inputs(
    payout_session_factory,
):  # B-52
    """End-to-end shape of the deadlock this fixes: the pool holds one UTXO per bet,
    so a round past MAX_TX_INPUTS participants could not be paid at all — the build
    failed with too_many_inputs, the round stayed "paying_out" retrying every 60s,
    and no new round could ever open behind it. It must now broadcast normally."""
    await _seed_paying_out_round(payout_session_factory)

    participants = MAX_TX_INPUTS + 1
    bet_sats = _POOL_AMOUNT_SATS // participants
    entries = [
        {"tx_hash": f"{i:064x}", "tx_pos": 0, "height": 10, "value": bet_sats}
        for i in range(participants)
    ]
    # The pool's total must cover the round's recorded pool_amount_sats, exactly as
    # on-chain: integer division above leaves a remainder, so top the last one up.
    entries[-1]["value"] += _POOL_AMOUNT_SATS - bet_sats * participants
    client = FakePayoutClient(entries)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client))

    await scheduler._trigger_payout(1)

    assert client.broadcasted
    async with payout_session_factory() as session:
        pending = (await session.scalars(select(PendingTransaction))).one()
        assert pending.status == "pending"
        assert len(Transaction.parse(bytes.fromhex(pending.raw_tx_hex)).vin) == participants

        events = [e.event_type for e in (await session.scalars(select(AuditLog))).all()]
        assert "payout_sent" in events
        assert "payout_failed" not in events


async def test_trigger_payout_broadcast_failure_leaves_a_recoverable_row(payout_session_factory):
    """Before B-25, a broadcast rejection here left nothing behind — no payout_txid,
    no PendingTransaction — because everything was persisted only after the
    broadcast. Now the intent is already durable, so the reconciler has something to
    resolve instead of the round being stuck with zero trace of what was attempted."""
    await _seed_paying_out_round(payout_session_factory)
    entries = [{"tx_hash": "44" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS + 100_000}]
    client = FakePayoutClient(entries, fail_broadcast=True)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client))

    await scheduler._trigger_payout(1)

    assert not client.broadcasted
    async with payout_session_factory() as session:
        round_ = await session.get(Round, 1)
        assert round_.payout_txid is not None  # durable, even though the broadcast failed

        pending = (await session.scalars(select(PendingTransaction))).one()
        assert pending.kind == "payout"
        assert pending.status == "building"  # not lost — the reconciler resolves this
        assert pending.current_txid == round_.payout_txid

        events = [e.event_type for e in (await session.scalars(select(AuditLog))).all()]
        assert "payout_failed" in events


async def test_trigger_payout_skips_when_already_in_flight(payout_session_factory):
    """A second call for a round that already has a non-terminal payout
    PendingTransaction must not build (and broadcast) another one — that would pay
    the winner twice."""
    winner_id = await _seed_paying_out_round(payout_session_factory)
    async with payout_session_factory() as session:
        round_ = await session.get(Round, 1)
        round_.payout_txid = "already-sent-txid"
        session.add(
            PendingTransaction(
                kind="payout",
                round_id=1,
                current_txid="already-sent-txid",
                fee_rate_sat_vb=1,
                raw_tx_hex="00",
                status="pending",
            )
        )
        await session.commit()

    entries = [{"tx_hash": "55" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS + 100_000}]
    client = FakePayoutClient(entries)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client))

    await scheduler._trigger_payout(1)

    assert not client.broadcasted
    async with payout_session_factory() as session:
        assert (await session.scalars(select(PendingTransaction))).all()  # still just the one seeded
        rows = (await session.scalars(select(PendingTransaction))).all()
        assert len(rows) == 1
        assert rows[0].current_txid == "already-sent-txid"


async def test_reserved_payout_outpoints_excludes_utxos_claimed_by_a_stale_payout(payout_session_factory):
    """A payout still "building"/"pending" for some round — most plausibly a stale
    one the reconciler hasn't abandoned yet — must keep its inputs off the table for
    a fresh payout attempt, or the same pool coins could be spent twice."""
    from embit import script
    from embit.transaction import Transaction, TransactionInput, TransactionOutput

    reserved_txid = "66" * 32
    raw_tx = (
        Transaction(
            vin=[TransactionInput(bytes.fromhex(reserved_txid), 2)],
            vout=[TransactionOutput(1_000_000, script.Script.from_address(_WINNER_ADDRESS))],
        )
        .serialize()
        .hex()
    )
    async with payout_session_factory() as session:
        session.add(
            PendingTransaction(
                kind="payout",
                round_id=99,
                current_txid="stale-payout-txid",
                fee_rate_sat_vb=1,
                raw_tx_hex=raw_tx,
                status="building",
            )
        )
        await session.commit()

        reserved = await _reserved_payout_outpoints(session)

    assert reserved == {(reserved_txid, 2)}


# --- B-26: a "paying_out" round must retry its payout automatically ---------------


async def test_trigger_payout_logs_a_failure_when_not_connected(payout_session_factory):
    """Before B-26, this early return logged nothing beyond a log line — invisible
    in /admin and unusable as a signal for an automatic retry."""
    await _seed_paying_out_round(payout_session_factory)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client=None))

    await scheduler._trigger_payout(1)

    async with payout_session_factory() as session:
        entries = (await session.scalars(select(AuditLog).where(AuditLog.event_type == "payout_failed"))).all()
        assert len(entries) == 1
        assert entries[0].payload_json.count("electrum client not connected") == 1


async def test_trigger_payout_logs_a_failure_when_fee_address_missing(payout_session_factory):
    winner_id = await _seed_paying_out_round(payout_session_factory)
    async with payout_session_factory() as session:
        config = (await session.scalars(select(RoundConfig))).one()
        config.fee_address = ""
        await session.commit()

    entries = [{"tx_hash": "77" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS + 100_000}]
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(FakePayoutClient(entries)))

    await scheduler._trigger_payout(1)

    async with payout_session_factory() as session:
        entry = (await session.scalars(select(AuditLog).where(AuditLog.event_type == "payout_failed"))).one()
        assert "no fee_address configured" in entry.payload_json
        assert entry.user_id == winner_id


async def test_tick_retries_a_stuck_paying_out_round_with_no_recent_failure(payout_session_factory):
    """The scenario B-26 exists for: a round stuck in "paying_out" (a prior failure,
    or a process restart mid-payout) with no non-terminal PendingTransaction. A
    fresh tick must retry rather than leaving it wedged forever."""
    await _seed_paying_out_round(payout_session_factory)
    entries = [{"tx_hash": "88" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS + 100_000}]
    client = FakePayoutClient(entries)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client))

    await scheduler._tick()

    assert client.broadcasted
    async with payout_session_factory() as session:
        round_ = await session.get(Round, 1)
        assert round_.payout_txid is not None
        pending = (await session.scalars(select(PendingTransaction))).one()
        assert pending.status == "pending"


async def test_tick_throttles_retry_after_a_recent_payout_failure(payout_session_factory):
    """A payout that just failed must not be retried on the very next tick, or a
    persistently-broken payout (e.g. no fee_address) would spam a retry — and a
    fresh payout_failed audit entry — every _TICK_INTERVAL_SECONDS."""
    await _seed_paying_out_round(payout_session_factory)
    async with payout_session_factory() as session:
        session.add(
            AuditLog(
                event_type="payout_failed",
                payload_json='{"round_id": 1, "reason": "insufficient pool UTXOs"}',
                round_id=1,
                created_at=datetime.now(timezone.utc),
            )
        )
        await session.commit()

    entries = [{"tx_hash": "99" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS + 100_000}]
    client = FakePayoutClient(entries)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client))

    await scheduler._tick()

    assert not client.broadcasted
    async with payout_session_factory() as session:
        assert (await session.scalars(select(PendingTransaction))).all() == []


async def test_tick_retries_once_the_throttle_window_has_elapsed(payout_session_factory):
    await _seed_paying_out_round(payout_session_factory)
    async with payout_session_factory() as session:
        session.add(
            AuditLog(
                event_type="payout_failed",
                payload_json='{"round_id": 1, "reason": "insufficient pool UTXOs"}',
                round_id=1,
                created_at=datetime.now(timezone.utc) - timedelta(seconds=120),
            )
        )
        await session.commit()

    entries = [{"tx_hash": "aa" * 32, "tx_pos": 0, "height": 10, "value": _POOL_AMOUNT_SATS + 100_000}]
    client = FakePayoutClient(entries)
    scheduler = RoundScheduler(payout_session_factory, FakePayoutListener(client))

    await scheduler._tick()

    assert client.broadcasted
    async with payout_session_factory() as session:
        pending = (await session.scalars(select(PendingTransaction))).one()
        assert pending.status == "pending"


# --- B-28: the draw must not seed itself from an uncorroborated header -----------


class CorroboratingListener:
    """A fake listener whose tip advances the moment a corroboration attempt
    fails, simulating a further block arriving — lets tests drive
    _wait_for_next_block's retry loop deterministically without real sleeps."""

    def __init__(self, *, responses: dict[int, bool], advance_to: dict[int, tuple[int, str]] | None = None):
        self.tip_height, self.tip_header_hex = next(iter(responses)), "aa"
        self._responses = dict(responses)
        self._advance_to = advance_to or {}
        self.corroboration_calls: list[int] = []

    async def corroborate_header(self, height: int, expected_hash: str) -> bool:
        self.corroboration_calls.append(height)
        result = self._responses[height]
        if not result and height in self._advance_to:
            self.tip_height, self.tip_header_hex = self._advance_to[height]
        return result


async def test_wait_for_next_block_accepts_an_immediately_corroborated_block(session_factory):
    listener = CorroboratingListener(responses={101: True})
    scheduler = RoundScheduler(session_factory, listener)

    height, block_hash = await scheduler._wait_for_next_block(
        round_id=1, tip_at_close=100, waiting_since=datetime.now(timezone.utc)
    )

    assert height == 101
    assert listener.corroboration_calls == [101]


async def test_wait_for_next_block_retries_after_a_failed_corroboration(session_factory):
    """B-28: an uncorroborated header must never be used — the wait keeps going
    until a later block's header *is* corroborated, logging why each time."""
    listener = CorroboratingListener(
        responses={101: False, 102: True}, advance_to={101: (102, "bb")}
    )
    scheduler = RoundScheduler(session_factory, listener)

    height, block_hash = await scheduler._wait_for_next_block(
        round_id=1, tip_at_close=100, waiting_since=datetime.now(timezone.utc)
    )

    assert height == 102
    assert listener.corroboration_calls == [101, 102]

    async with session_factory() as session:
        events = [e.event_type for e in (await session.scalars(select(AuditLog))).all()]
        assert events == ["draw_header_corroboration_failed"]


# --- B-36: a stalled draw must be visible, not a silent frozen wait --------------


class StallingListener:
    """A tip that never advances until the test decides it should — used to drive
    _wait_for_next_block's stall-detection past _DRAW_STALL_THRESHOLD_SECONDS
    without a real 6-minute wait."""

    def __init__(self):
        self.tip_height = 100
        self.tip_header_hex = None

    async def corroborate_header(self, height: int, expected_hash: str) -> bool:
        return True


async def test_wait_for_next_block_logs_a_stall_audit_entry_past_the_threshold(session_factory, monkeypatch):
    import app.rounds.scheduler as scheduler_module

    listener = StallingListener()
    scheduler = RoundScheduler(session_factory, listener)
    start = datetime.now(timezone.utc)

    class _FakeClock:
        now = start

    def fake_now(tz=None):
        return _FakeClock.now

    async def fake_sleep(seconds: float) -> None:
        _FakeClock.now += timedelta(seconds=seconds)
        # Past the stall threshold, but before it would repeat: unblock the wait
        # by making a (corroborated) block appear, so the test terminates.
        if _FakeClock.now >= start + timedelta(seconds=scheduler_module._DRAW_STALL_THRESHOLD_SECONDS + 30):
            listener.tip_height = 101
            listener.tip_header_hex = "aa"

    monkeypatch.setattr(scheduler_module, "datetime", type("_D", (), {"now": staticmethod(fake_now)}))
    monkeypatch.setattr(scheduler_module.asyncio, "sleep", fake_sleep)

    height, block_hash = await scheduler._wait_for_next_block(round_id=1, tip_at_close=100, waiting_since=start)

    assert height == 101

    async with session_factory() as session:
        entries = (await session.scalars(select(AuditLog).where(AuditLog.event_type == "draw_stalled"))).all()
        assert len(entries) == 1
        assert entries[0].round_id == 1