webmail-client/backend/tests/sync/test_acceptance.py

"""Sync-engine v2 acceptance tests 1-3 (refactor §10).

These drive the real :class:`SyncCoordinator` against the test Redis with the
per-account IMAP loops replaced by fakes, so the lease protocol — single
ownership, crash failover, and create/delete lifecycle — is exercised without a
live mail server. Intervals are shortened so the tests run in seconds.

  1. Two workers, N accounts: every account owned by exactly one worker, exactly
     one live (IDLE) connection each.
  2. Kill a worker: its orphaned accounts are re-owned by the survivor once the
     leases expire, still one connection each (no double ownership).
  3. A newly-live account is claimed without a restart; a removed account's
     connection is torn down (lease fencing), with no further ownership.
"""

from __future__ import annotations

import asyncio
import uuid
from collections import defaultdict
from collections.abc import AsyncIterator, Callable
from datetime import UTC, datetime, timedelta

import pytest_asyncio
from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool

from app.config import get_settings
from app.db import session_scope
from app.events import get_event_bus
from app.identity.passwords import hash_password
from app.models import Account, Mailbox, Message, User
from app.redis_client import get_redis
from app.sync import coordinator as coordinator_mod
from app.sync import leases, runner
from app.sync.coordinator import SyncCoordinator

# --- fakes: replace the real IMAP loops with connection-counting stand-ins ----


@pytest_asyncio.fixture
async def connections(monkeypatch) -> AsyncIterator[dict[uuid.UUID, int]]:  # type: ignore[no-untyped-def]
    """Patch the runner's loops; yield a live (hot-tier IDLE) connection count."""
    conns: dict[uuid.UUID, int] = defaultdict(int)

    async def fake_live_loop(account_id: uuid.UUID, folder: str, stop: asyncio.Event) -> None:
        conns[account_id] += 1  # one persistent IDLE connection while alive (hot)
        try:
            await stop.wait()
        finally:
            conns[account_id] -= 1

    async def fake_secondary(
        account_id: uuid.UUID, stop: asyncio.Event | None = None, interval: float | None = None
    ) -> None:
        assert stop is not None
        await stop.wait()

    async def fake_poll_all(
        account_id: uuid.UUID, interval: float, stop: asyncio.Event | None = None
    ) -> None:
        assert stop is not None  # warm/cold: no persistent connection, just wait
        await stop.wait()

    monkeypatch.setattr(runner, "live_loop", fake_live_loop)
    monkeypatch.setattr(runner, "secondary_poll_loop", fake_secondary)
    monkeypatch.setattr(runner, "poll_all_loop", fake_poll_all)
    yield conns


@pytest_asyncio.fixture
async def make_accounts() -> AsyncIterator[Callable[[int, str], object]]:
    """Factory creating N live accounts; cleans up its users afterwards."""
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    sm = async_sessionmaker(engine, expire_on_commit=False)
    user_ids: list[uuid.UUID] = []

    async def _make(n: int, state: str = "live") -> list[uuid.UUID]:
        async with sm() as session:
            user = User(
                handle=f"acc{uuid.uuid4().hex[:8]}@test.local", password_hash=hash_password("pw")
            )
            session.add(user)
            await session.flush()
            user_ids.append(user.id)
            ids: list[uuid.UUID] = []
            for _ in range(n):
                account = Account(
                    user_id=user.id,
                    email=f"{uuid.uuid4().hex[:10]}@example.com",
                    display_name="Acct",
                    imap_config={},
                    smtp_config={},
                    secret_ct=b"",
                    discovery_src="manual",
                    state=state,
                )
                session.add(account)
                await session.flush()
                ids.append(account.id)
            await session.commit()
            return ids

    try:
        yield _make
    finally:
        async with sm() as session:
            for uid in user_ids:
                await session.execute(delete(User).where(User.id == uid))
            await session.commit()
        await engine.dispose()


async def _no_reconcile() -> None:
    """Stub: keep the shared dev DB's real live accounts out of the test."""
    return None


def _fast(c: SyncCoordinator) -> SyncCoordinator:
    """Compress cadences and drive sync:live explicitly (no DB reconcile).

    The tests share the dev Postgres, which holds real live accounts; reconciling
    from it would pull them into the test. We seed ``sync:live`` ourselves instead.
    """
    c.claim_interval = 0.05
    c.claim_jitter = 0.0
    c.lease_renew = 0.3
    c.lease_ttl = 2  # short so a crashed worker's leases expire quickly
    c.reconcile_interval = 1000.0
    c.heartbeat_refresh = 1000.0
    c.rebalance_interval = 1000.0
    c.tier_eval_interval = 1000.0  # lease tests pin the tier; no eval churn
    c._reconcile_once = _no_reconcile  # type: ignore[method-assign]
    return c


async def _seed(ids: set[uuid.UUID], tier: str = "hot") -> None:
    """Register accounts as live and pin their tier (hot => persistent IDLE)."""
    redis = get_redis()
    for account_id in ids:
        await leases.mark_live(redis, account_id)
        await redis.set(leases.tier_key(account_id), tier)


async def _user_of(account_id: uuid.UUID) -> uuid.UUID:
    async with session_scope() as session:
        user_id = await session.scalar(select(Account.user_id).where(Account.id == account_id))
    assert user_id is not None
    return user_id


async def _add_recent_mail(account_id: uuid.UUID, when: datetime) -> None:
    """Insert one INBOX message dated ``when`` so the account reads as active."""
    user_id = await _user_of(account_id)
    async with session_scope() as session:
        mailbox = Mailbox(account_id=account_id, user_id=user_id, path="INBOX", special_use="inbox")
        session.add(mailbox)
        await session.flush()
        session.add(
            Message(
                user_id=user_id,
                account_id=account_id,
                mailbox_id=mailbox.id,
                uid=1,
                internal_date=when,
                blob_sha256="0" * 64,
            )
        )
        await session.commit()


async def _until(pred: Callable[[], bool], timeout_s: float = 6.0) -> None:
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout_s
    while loop.time() < deadline:
        if pred():
            return
        await asyncio.sleep(0.05)
    assert pred(), "condition not met within timeout"


# --- acceptance 1: single ownership, one connection per account ---------------


async def test_acceptance_1_single_ownership(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    ids = set(await make_accounts(20))
    await _seed(ids)
    c1 = _fast(SyncCoordinator("w1"))
    c2 = _fast(SyncCoordinator("w2"))
    c1.rebalance_interval = c2.rebalance_interval = 0.3  # let the load even out
    # Start concurrently so the two workers race for the leases.
    await asyncio.gather(c1.start(), c2.start())
    try:
        # Both workers end up holding accounts (via the initial race + rebalance),
        # together covering all 20 with no account owned twice.
        await _until(lambda: c1.owned | c2.owned == ids and bool(c1.owned and c2.owned))
        assert c1.owned & c2.owned == set()
        # At steady state each account has exactly one live IDLE connection.
        await _until(lambda: all(connections[a] == 1 for a in ids))
    finally:
        await c1.stop()
        await c2.stop()
    await _until(lambda: all(connections[a] == 0 for a in ids))


# --- acceptance 2: crash failover --------------------------------------------


async def test_acceptance_2_failover(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    ids = set(await make_accounts(10))
    await _seed(ids)
    c1 = _fast(SyncCoordinator("w1"))
    await c1.start()
    try:
        await _until(lambda: c1.owned == ids)  # sole worker owns everything

        c2 = _fast(SyncCoordinator("w2"))
        await c2.start()

        # Simulate `kill -9` on w1: connections drop and renewals stop, but the
        # leases are NOT released — they must expire before w2 can take over.
        for task in c1._tasks:
            task.cancel()
        await asyncio.gather(*c1._tasks, return_exceptions=True)
        c1._tasks = []
        for r in list(c1._runners.values()):
            await r.stop()
        c1._runners.clear()

        try:
            # w2 re-owns every orphan once the (short) leases expire.
            await _until(lambda: c2.owned == ids, timeout_s=8.0)
            assert all(connections[a] == 1 for a in ids)  # one connection each, no dupes
        finally:
            await c2.stop()
    finally:
        # w1 already crashed; clear any residual leases it left behind.
        await get_redis().flushdb()
    await _until(lambda: all(connections[a] == 0 for a in ids))


# --- acceptance 3: create / delete lifecycle without restart ------------------


async def test_acceptance_3_create_delete(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    c = _fast(SyncCoordinator("w1"))
    await c.start()  # starts with no live accounts
    try:
        assert c.owned == set()

        # A newly-live account (as the bootstrap job's _set_state('live') would do):
        # register it in sync:live; the running coordinator must claim it, no restart.
        (aid,) = await make_accounts(1)
        await get_redis().set(leases.tier_key(aid), "hot")  # persistent IDLE
        await leases.mark_live(get_redis(), aid)
        await _until(lambda: aid in c.owned)
        assert connections[aid] == 1

        # Removing it (as delete_account does) fences the runner: lease gone ->
        # next renewal fails -> connection torn down, no further ownership.
        await leases.mark_not_live(get_redis(), aid)
        await _until(lambda: aid not in c.owned)
        await _until(lambda: connections[aid] == 0)
    finally:
        await c.stop()


async def test_event_driven_teardown(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    """§3: an account.state_changed event tears the runner down at once.

    Leases are kept long (no fencing) and the verify loop slow, so the only path
    that can stop the runner in time is the cross-process event subscriber.
    """
    (aid,) = await make_accounts(1)
    user_id = uuid.uuid4()
    await _seed({aid})
    c = _fast(SyncCoordinator("w1"))
    c.lease_ttl = 60  # lease stays valid: fencing cannot explain the teardown
    c.lease_renew = 30
    await c.start()
    try:
        await _until(lambda: aid in c.owned)
        assert connections[aid] == 1

        # As delete_account does: drop it from sync:live (so it isn't re-claimed),
        # then emit the event. The lease is left intact, so only the event — not
        # fencing — can explain the teardown.
        await get_redis().srem(leases.LIVE_SET, str(aid))
        await get_event_bus().publish(
            "account.state_changed",
            {"user_id": str(user_id), "account_id": str(aid), "state": "deleted"},
        )
        await _until(lambda: aid not in c.owned, timeout_s=3.0)
        await _until(lambda: connections[aid] == 0, timeout_s=3.0)
    finally:
        await c.stop()


# --- §5 activity tiers -------------------------------------------------------


async def test_tier_session_promotes_to_hot(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    (aid,) = await make_accounts(1)
    user_id = await _user_of(aid)
    await leases.mark_live(get_redis(), aid)  # claimed as warm by default
    await leases.mark_user_active(get_redis(), user_id, 300)  # active session -> hot
    c = _fast(SyncCoordinator("w1"))
    c.tier_eval_interval = 0.1
    await c.start()
    try:
        await _until(lambda: aid in c.owned)
        await _until(lambda: connections[aid] == 1)  # hot => persistent IDLE
        assert c._runners[aid].tier == "hot"
    finally:
        await c.stop()


async def test_tier_no_session_no_mail_goes_cold(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    (aid,) = await make_accounts(1)
    await leases.mark_live(get_redis(), aid)  # no session, no mail
    c = _fast(SyncCoordinator("w1"))
    c.tier_eval_interval = 0.1
    await c.start()
    try:
        await _until(lambda: aid in c.owned)
        await _until(lambda: c._runners[aid].tier == "cold")
        assert connections[aid] == 0  # cold holds no persistent connection
    finally:
        await c.stop()


async def test_tier_recent_mail_promotes_to_hot(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    (aid,) = await make_accounts(1)
    await _add_recent_mail(aid, datetime.now(UTC))  # activity within TIER_HOT_ACTIVITY_H
    await leases.mark_live(get_redis(), aid)
    c = _fast(SyncCoordinator("w1"))
    c.tier_eval_interval = 0.1
    await c.start()
    try:
        await _until(lambda: aid in c.owned)
        await _until(lambda: connections[aid] == 1)  # recent mail => hot
        assert c._runners[aid].tier == "hot"
    finally:
        await c.stop()


async def test_user_active_event_promotes_immediately(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    aid0, aid1 = await make_accounts(2)  # both under one user
    user_id = await _user_of(aid0)
    await leases.mark_live(get_redis(), aid0)
    await leases.mark_live(get_redis(), aid1)
    c = _fast(SyncCoordinator("w1"))  # tier eval disabled: only the event can promote
    await c.start()
    try:
        await _until(lambda: {aid0, aid1} <= c.owned)
        assert connections[aid0] == 0 and connections[aid1] == 0  # warm to start
        await get_event_bus().publish("user.active", {"user_id": str(user_id)})
        await _until(lambda: connections[aid0] == 1 and connections[aid1] == 1, timeout_s=3.0)
    finally:
        await c.stop()


async def test_hot_budget_demotes_least_active(connections, make_accounts) -> None:  # type: ignore[no-untyped-def]
    aid_old, aid_new = await make_accounts(2)  # activity-hot, same user, NO session
    now = datetime.now(UTC)
    await _add_recent_mail(aid_old, now - timedelta(hours=1))  # least recently active
    await _add_recent_mail(aid_new, now)
    await leases.mark_live(get_redis(), aid_old)
    await leases.mark_live(get_redis(), aid_new)
    c = _fast(SyncCoordinator("w1"))
    c.tier_eval_interval = 0.1
    c.hot_budget = 1  # only one may be hot
    await c.start()
    try:
        await _until(lambda: {aid_old, aid_new} <= c.owned)
        # The newer account stays hot; the older is shed to warm by the budget.
        await _until(lambda: c._runners[aid_new].tier == "hot")
        await _until(lambda: c._runners[aid_old].tier == "warm")
        assert connections[aid_new] == 1 and connections[aid_old] == 0
    finally:
        await c.stop()


def test_worker_id_is_unique() -> None:
    assert coordinator_mod._new_worker_id() != coordinator_mod._new_worker_id()