"""Tests for grpc_client.Subscription: apply throughput, batch-drain, shutdown,
and error-handling behaviors.

The subscription reader thread decodes ops and enqueues them on a bounded queue;
a single applier coroutine batch-drains the queue and applies each op with a
plain await, so the cross-thread hand-off is amortized over a whole batch rather
than paid once per op.  These tests drive the REAL apply path
(``ref_handler`` + ``handle_ref_update2``) over a synthetic in-memory op stream
in the production topology (a long-lived event loop in its own thread plus a
separate reader thread).  No network / live server / credentials.

Organized by concern into test classes:
  - TestSubscriptionApplyThroughput  -- steady-state apply throughput plus the
    deterministic ops-per-hand-off structural guard.
  - TestSubscriptionShutdown         -- close()/teardown under backpressure and
    out-of-band (sentinel-free) applier termination.
  - TestDrainBatchCap                -- the per-event-loop-turn batch cap and the
    reader-done drain termination condition.
  - TestApplyErrorThrottling         -- callback-apply failures are skipped and
    the traceback is logged throttled, not per-op.
  - TestRunDeterministicShutdown     -- run() does not orphan a sibling when the
    reader or the applier raises, and surfaces the failure.

Throughput characterization (60,000 steady-state ops, identical real mutation):
a per-op cross-thread round-trip applier measures ~42,000 ops/s; the batch-drain
applier measures ~570,000-600,000 ops/s (~13x).  The batch applies up to
_MAX_DRAIN_BATCH (1024) ops per event-loop hand-off, vs one per op for a per-op
applier.

THROUGHPUT_FLOOR_OPS_PER_S is 120,000 ops/s: ~3.3x above the per-op baseline (so
a per-op applier fails this test) and ~4x below the measured batch-drain number
(so normal run-to-run variance does not flake).  ``test_per_op_baseline_is_
below_floor`` characterizes the baseline so the floor's separation from it is
self-checking.

The always-on guard is ``test_batch_drain_amortizes_event_loop_handoffs`` -- a
deterministic structural assertion (many ops per event-loop hand-off) with no
wall-clock dependence, so it cannot flake on slow/noisy CI.  The wall-clock
floor case ``test_batch_drain_throughput_above_floor`` is secondary and is
skipped under CI (where the floor, calibrated on a dev box, can flake on slower
shared runners) or when PPA2125_SKIP_THROUGHPUT_WALLCLOCK is set -- without
losing the structural coverage, and without lowering the floor for local runs.
"""
import asyncio
import concurrent.futures
import os
import threading
import time
import unittest

import grpc

from preseem.grpc_client import (Subscription, _SUBSCRIPTION_QUEUE_MAXSIZE,
                                 _MAX_DRAIN_BATCH)
from preseem.handoff_queue import HandoffQueue
from preseem import PreseemNetworkMetadataModel, NetworkMetadataIndex
from preseem_grpc_model import network_metadata_pb2

RA = network_metadata_pb2.NetworkMetadataReferenceEntityAction
AT = network_metadata_pb2.NetworkMetadataAttribute


def _q_full(hq):
    """White-box: is the HandoffQueue at capacity? (no public size accessor)."""
    with hq._cond:
        return len(hq._items) >= hq._maxsize


# Mirrors the netmeta-agent production index config.
NM_INDEXES = [
    NetworkMetadataIndex("snmp_session", "dev_mac"),
    NetworkMetadataIndex("session", "service"),
]

# Steady-state (post-LOADED) op count to time -- large enough for a stable
# measurement and to exercise many drain batches.
SOAK_OPS = 60_000
# Snapshot establishes a populated model so steady-state SETs hit the same
# index-update / ref-lookup paths as production.
SNAPSHOT_OPS = 20_000

# See module docstring for derivation.
THROUGHPUT_FLOOR_OPS_PER_S = 120_000

# The wall-clock floor was calibrated on a dev box; shared CI runners are slower
# and noisier, so the floor case is skipped under CI (GitHub Actions and most CI
# set CI=true) or when explicitly requested.  The deterministic structural guard
# (test_batch_drain_amortizes_event_loop_handoffs) has no wall-clock dependence
# and always runs, so the regression coverage is preserved regardless.
_SKIP_WALLCLOCK = bool(
    os.environ.get("PPA2125_SKIP_THROUGHPUT_WALLCLOCK") or os.environ.get("CI"))


def _set_attr(attrs, name, *, s=None, i=None):
    a = attrs[name]
    a.name = name
    if s is not None:
        a.type = AT.STRING
        a.stringVal = s
    elif i is not None:
        a.type = AT.INT64
        a.int64Val = i


def _make_ref_action(action, n, ref_type):
    m = RA()
    m.action = action
    if action != RA.SET:
        return m
    e = m.entity
    e.type = ref_type
    e.value = "{}-{}".format(ref_type, n)
    if ref_type == "session":
        _set_attr(e.attributes, "service", s="service-{}".format(n % 5000))
        _set_attr(e.attributes, "plan", s=("gold" if n % 3 else "silver"))
        _set_attr(e.attributes, "down_kbps", i=100000 + (n % 50) * 1000)
    elif ref_type == "snmp_session":
        _set_attr(e.attributes, "dev_mac", s="02:%010x" % n)
    else:
        _set_attr(e.attributes, "name", s="{}-name-{}".format(ref_type, n))
    return m


def _build_stream(snapshot, soak):
    """A full subscription stream: `snapshot` SETs, a LOADED marker, then `soak`
    steady-state SETs.  ref_handler batches the snapshot into one Op.SYN and then
    yields one Op.SET per soak message (the steady-state apply path under test).
    """
    ref_type = ("session", "snmp_session", "device")
    msgs = [_make_ref_action(RA.SET, n, ref_type[n % 3]) for n in range(snapshot)]
    msgs.append(_make_ref_action(RA.LOADED, 0, ""))
    for j in range(soak):
        n = (j % snapshot) if (j % 4) else (snapshot + j)
        msgs.append(_make_ref_action(RA.SET, n, ref_type[n % 3]))
    return msgs


class _StreamIterator(object):
    """Synthetic gRPC iterator: yields the messages via .next() then raises a
    CANCELLED RpcError so the reader/handler stops cleanly (no live server)."""

    def __init__(self, msgs):
        self._it = iter(msgs)

    def cancel(self):
        pass

    def next(self):
        try:
            return next(self._it)
        except StopIteration:
            raise _CancelledRpcError()


class _CancelledRpcError(grpc.RpcError):

    def code(self):
        return grpc.StatusCode.CANCELLED

    def details(self):
        return 'Locally cancelled by application!'


def _loop_in_thread():
    """Start an event loop in its own thread -- the production topology (a
    long-lived loop thread, separate from the subscription reader thread)."""
    loop = asyncio.new_event_loop()
    ready = threading.Event()

    def _runner():
        asyncio.set_event_loop(loop)
        loop.call_soon(ready.set)
        loop.run_forever()

    threading.Thread(target=_runner, daemon=True).start()
    ready.wait()
    return loop


def _fresh_model():
    return PreseemNetworkMetadataModel(client=None,
                                       company_id=None,
                                       refs_by_type=True,
                                       indexes=NM_INDEXES)


def _extract_ops(model, msgs):
    """Drive the real ref_handler to collect (op, payload) tuples OUTSIDE the
    timed region.  Returns [Op.SYN(snapshot), Op.SET, Op.SET, ...]."""
    ops = []
    try:
        for op in model.ref_handler(_StreamIterator(msgs)):
            ops.append(op)
    except grpc.RpcError:
        pass
    return ops


def _new_subscription(model, loop):
    """Build a real production Subscription bound to `loop`, with no live client
    or connection (we drive its applier directly over an in-memory queue)."""
    sub = Subscription.__new__(Subscription)
    sub._client = type("_C", (), {"host": "synthetic"})()
    sub._op = None
    sub._req = None
    sub._cbk = model.handle_ref_update2
    sub._handler = model.ref_handler
    sub._executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
    sub._loop = loop
    sub._iter = None
    sub._queue = HandoffQueue(_SUBSCRIPTION_QUEUE_MAXSIZE)
    sub._closing = threading.Event()
    sub._apply_error_last_log = None
    sub._apply_errors_since_log = 0
    sub._kind = "test"
    return sub


def _feed(sub, ops):
    """Mirror the production reader: enqueue ops, then mark the hand-off queue
    producer_done (out-of-band, no in-band sentinel) so _apply terminates."""
    for op in ops:
        sub._queue.put(op)
    sub._queue.producer_done()


def _apply_via_production_batch_drain(model, ops, loop):
    """Apply `ops` through the REAL production applier (Subscription._apply +
    _drain_batch): a reader thread enqueues, the applier batch-drains on `loop`.
    Op.SYN (ops[0]) is applied untimed; the steady Op.SET region is timed.
    Returns (n_timed, elapsed, drain_handoffs)."""
    sub = _new_subscription(model, loop)
    handoffs = {"n": 0}
    real_drain = sub._drain_batch

    def counting_drain():
        batch, done = real_drain()
        if batch:  # count only productive hand-offs
            handoffs["n"] += 1
        return batch, done

    sub._drain_batch = counting_drain

    # Phase 1: untimed SYN prefix through the same applier.
    applier = asyncio.run_coroutine_threadsafe(sub._apply(), loop)
    threading.Thread(target=_feed, args=(sub, ops[:1])).start()
    applier.result()

    # Phase 2: timed steady region.  A fresh hand-off queue (phase 1's was marked
    # producer_done/stopped) plus a cleared closing flag drive a fresh stream.
    sub._queue = HandoffQueue(_SUBSCRIPTION_QUEUE_MAXSIZE)
    sub._closing.clear()
    handoffs["n"] = 0
    timed = ops[1:]
    applier = asyncio.run_coroutine_threadsafe(sub._apply(), loop)
    t0 = time.perf_counter()
    threading.Thread(target=_feed, args=(sub, timed)).start()
    applier.result()
    elapsed = time.perf_counter() - t0
    return len(timed), elapsed, handoffs["n"]


def _apply_via_per_op(model, ops, loop):
    """Apply `ops` the OLD way: a reader thread does one
    run_coroutine_threadsafe(...).result() per op against `loop` (the topology of
    the pre-fix Subscription._run).  SYN untimed, steady region timed.
    Returns (n_timed, elapsed)."""
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
    cbk = model.handle_ref_update2

    def dispatch(these):
        for op in these:
            asyncio.run_coroutine_threadsafe(cbk(*op), loop).result()

    async def run(these):
        await loop.run_in_executor(executor, dispatch, these)

    asyncio.run_coroutine_threadsafe(run(ops[:1]), loop).result()  # SYN, untimed
    timed = ops[1:]
    t0 = time.perf_counter()
    asyncio.run_coroutine_threadsafe(run(timed), loop).result()
    return len(timed), time.perf_counter() - t0


class TestSubscriptionApplyThroughput(unittest.TestCase):
    """Verifies the batch-drain throughput and per-hand-off structure of the
    Subscription apply path."""

    def setUp(self):
        self._loop = _loop_in_thread()
        msgs = _build_stream(SNAPSHOT_OPS, SOAK_OPS)
        # Extract ops once, outside any timed region, so we measure apply, not
        # handler decode.  Each arm gets its own fresh model + op list.
        self._make_ops = lambda: (lambda m: (m, _extract_ops(m, msgs)))(_fresh_model())

    def tearDown(self):
        self._loop.call_soon_threadsafe(self._loop.stop)

    @unittest.skipIf(
        _SKIP_WALLCLOCK, "wall-clock throughput floor skipped under CI / when "
        "PPA2125_SKIP_THROUGHPUT_WALLCLOCK is set; the deterministic structural "
        "guard still runs.")
    def test_batch_drain_throughput_above_floor(self):
        model, ops = self._make_ops()
        n, elapsed, handoffs = _apply_via_production_batch_drain(model, ops, self._loop)
        ops_per_s = n / elapsed
        self.assertGreaterEqual(
            ops_per_s, THROUGHPUT_FLOOR_OPS_PER_S,
            "production batch-drain apply throughput {:,.0f} ops/s ({:,} ops in "
            "{:.4f}s, {} drain hand-offs) fell below the floor of {:,} "
            "ops/s -- this usually means the applier regressed to a per-op "
            "cross-thread round-trip instead of batch-draining.".format(
                ops_per_s, n, elapsed, handoffs, THROUGHPUT_FLOOR_OPS_PER_S))

    def test_batch_drain_amortizes_event_loop_handoffs(self):
        """Deterministic structural guard (no wall-clock dependence): the applier
        must apply MANY ops per event-loop hand-off.  A per-op applier does one
        hand-off per op; the batch-drain applier drains the whole queue per turn."""
        model, ops = self._make_ops()
        n, _elapsed, handoffs = _apply_via_production_batch_drain(
            model, ops, self._loop)
        self.assertGreater(handoffs, 0)
        ops_per_handoff = n / handoffs
        self.assertGreater(
            ops_per_handoff, 10,
            "batch-drain applied only {:.1f} ops per event-loop hand-off "
            "({} ops, {} hand-offs); the fix should amortize many ops per "
            "hand-off (a regression to per-op gives ~1).".format(
                ops_per_handoff, n, handoffs))

    def test_per_op_baseline_is_below_floor(self):
        """Sanity check (not the primary guard): characterize the per-op
        baseline and confirm it is BELOW the floor, so the floor in
        test_batch_drain_throughput_above_floor would actually catch a regression
        back to per-op."""
        model, ops = self._make_ops()
        n, elapsed = _apply_via_per_op(model, ops, self._loop)
        ops_per_s = n / elapsed
        self.assertLess(
            ops_per_s, THROUGHPUT_FLOOR_OPS_PER_S,
            "per-op baseline measured {:,.0f} ops/s ({:,} ops in {:.4f}s), "
            "which is NOT below the floor {:,} -- on this host the floor would "
            "not distinguish a per-op regression, so the floor needs "
            "review.".format(ops_per_s, n, elapsed, THROUGHPUT_FLOOR_OPS_PER_S))


def _drain_batch_subscription(maxsize):
    """A bare Subscription with just the fields _drain_batch reads, so we can
    exercise the drain in isolation."""
    sub = Subscription.__new__(Subscription)
    sub._queue = HandoffQueue(maxsize)
    sub._closing = threading.Event()
    return sub


class TestSubscriptionShutdown(unittest.TestCase):
    """Verifies shutdown under backpressure: the reader thread never hangs on a
    full queue, close()/loop-teardown terminates it promptly, and applier
    termination is out-of-band -- it does not depend on an in-band sentinel that
    a concurrent drainer could swallow."""

    def _flooding_subscription(self, maxsize):
        """A Subscription whose handler yields ops forever and whose applier is
        never scheduled, so the bounded queue fills and the reader parks in
        _put -- the backpressure scenario."""
        sub = Subscription.__new__(Subscription)
        sub._client = type("_C", (), {"host": "synthetic"})()
        sub._op = None
        sub._req = None
        sub._cbk = None
        sub._iter = type("_It", (), {"cancel": lambda self: None})()

        def _forever(_iter):
            n = 0
            while True:
                yield ("SET", n)
                n += 1

        sub._handler = _forever
        sub._loop = asyncio.new_event_loop()
        sub._queue = HandoffQueue(maxsize)
        sub._closing = threading.Event()
        sub._kind = "test"
        return sub

    def test_close_during_iter_creation_does_not_block_reader(self):
        """TOCTOU: if close() lands while _iter is still None, _request_stop
        cannot cancel it; the reader must catch the closing flag via the
        post-creation re-check instead of blocking forever in next()."""
        sub = Subscription.__new__(Subscription)
        sub._client = type("_C", (), {"host": "synthetic"})()
        sub._op = None
        sub._req = None
        sub._cbk = None
        sub._handler = None
        sub._iter = None
        sub._loop = asyncio.new_event_loop()
        sub._queue = HandoffQueue(5)
        sub._closing = threading.Event()
        sub._kind = "test"

        class _BlockForeverIter(object):

            def next(self):
                # Would hang forever if the reader ever reached the read.
                threading.Event().wait()

            def cancel(self):
                pass

        def _create(_req):
            # Simulate close() racing in exactly while _iter is still None:
            # the flag is set before this returns, but _request_stop already saw
            # _iter is None and could not cancel.
            sub._closing.set()
            return _BlockForeverIter()

        sub._op = _create

        reader = threading.Thread(target=sub._run)
        reader.start()
        reader.join(timeout=3.0)
        self.assertFalse(
            reader.is_alive(),
            "reader blocked in next() despite close() during iter creation")

    def test_close_unwedges_reader_blocked_in_handler_next(self):
        """Handler path: the reader spends the whole connection inside the
        handler generator's `while True: it.next()`, which never re-checks
        _closing.  close() must still unwedge it via the iterator cancel."""
        sub = Subscription.__new__(Subscription)
        sub._client = type("_C", (), {"host": "synthetic"})()
        sub._req = None
        sub._cbk = None
        sub._iter = None
        sub._loop = asyncio.new_event_loop()
        sub._queue = HandoffQueue(5)
        sub._closing = threading.Event()
        sub._kind = "test"

        class _CancellableBlockingIter(object):

            def __init__(self):
                self._cancelled = threading.Event()
                self.parked = threading.Event()

            def next(self):
                self.parked.set()
                self._cancelled.wait()  # blocks until cancel()
                raise _CancelledRpcError()

            def cancel(self):
                self._cancelled.set()

        it = _CancellableBlockingIter()
        sub._op = lambda _req: it

        def _handler(iterator):
            while True:  # mimics ref_handler: never re-checks _closing
                yield iterator.next()

        sub._handler = _handler

        reader = threading.Thread(target=sub._run)
        reader.start()
        self.assertTrue(it.parked.wait(timeout=2.0))  # inside handler next()
        sub.close()
        reader.join(timeout=3.0)
        self.assertFalse(reader.is_alive(),
                         "reader stayed blocked in handler next() after close()")
        self.assertTrue(it._cancelled.is_set())

    def test_create_failure_retries_via_interruptible_backoff(self):
        """A single failed create attempt must route through _run's own
        _closing-aware backoff, so the reader keeps retrying but close()
        interrupts it promptly."""
        sub = Subscription.__new__(Subscription)
        sub._client = type("_C", (), {"host": "synthetic"})()
        sub._req = None
        sub._cbk = None
        sub._handler = None
        sub._iter = None
        sub._loop = asyncio.new_event_loop()
        sub._queue = HandoffQueue(5)
        sub._closing = threading.Event()
        sub._kind = "test"
        attempts = {"n": 0}

        def _always_fails(_req):
            attempts["n"] += 1
            raise RuntimeError("cannot connect")

        sub._op = _always_fails
        reader = threading.Thread(target=sub._run)
        reader.start()
        # Let it make at least one failed attempt and enter the backoff.
        deadline = time.time() + 2.0
        while attempts["n"] < 1 and time.time() < deadline:
            time.sleep(0.01)
        self.assertGreaterEqual(attempts["n"], 1)
        t0 = time.time()
        sub.close()
        reader.join(timeout=3.0)
        self.assertFalse(reader.is_alive())
        self.assertLess(time.time() - t0, 5.0)  # not the full 10s backoff

    def test_create_returning_none_does_not_deref(self):
        """If create returns None (e.g. retry exhaustion), the reader must not
        call None.next(); it backs off and retries instead."""
        sub = Subscription.__new__(Subscription)
        sub._client = type("_C", (), {"host": "synthetic"})()
        sub._req = None
        sub._cbk = None
        sub._handler = None
        sub._iter = None
        sub._loop = asyncio.new_event_loop()
        sub._queue = HandoffQueue(5)
        sub._closing = threading.Event()
        sub._kind = "test"
        calls = {"n": 0}

        def _returns_none(_req):
            calls["n"] += 1
            return None

        sub._op = _returns_none
        reader = threading.Thread(target=sub._run)
        reader.start()
        deadline = time.time() + 2.0
        while calls["n"] < 1 and time.time() < deadline:
            time.sleep(0.01)
        self.assertGreaterEqual(calls["n"], 1)
        # No AttributeError crashed the thread; close() stops it cleanly.
        sub.close()
        reader.join(timeout=3.0)
        self.assertFalse(reader.is_alive())

    def test_close_unblocks_reader_parked_on_full_queue(self):
        sub = self._flooding_subscription(maxsize=5)
        reader = threading.Thread(target=sub._run)
        reader.start()
        # Let the reader fill the queue and park in put().
        deadline = time.time() + 2.0
        while not _q_full(sub._queue) and time.time() < deadline:
            time.sleep(0.01)
        self.assertTrue(_q_full(sub._queue), "queue did not fill under flood")
        self.assertTrue(reader.is_alive())

        sub.close()
        reader.join(timeout=3.0)
        self.assertFalse(
            reader.is_alive(),
            "reader thread did not exit promptly after close() while parked on "
            "a full queue -- it would leak")
        # The reader signalled completion out-of-band (producer_done), not via a
        # queue item.
        self.assertTrue(sub._queue._producer_done)

    def test_reader_exits_without_draining_the_queue(self):
        """close() does NOT drain the queue (draining it could race the reader's
        out-of-band completion); the reader still exits, freed by the queue's
        stop() waking its parked put()."""
        sub = self._flooding_subscription(maxsize=5)
        reader = threading.Thread(target=sub._run)
        reader.start()
        deadline = time.time() + 2.0
        while not _q_full(sub._queue) and time.time() < deadline:
            time.sleep(0.01)
        sub.close()
        reader.join(timeout=3.0)
        self.assertFalse(reader.is_alive())
        # Queue contents are left intact by close(); nothing is consumed.
        self.assertTrue(_q_full(sub._queue))

    def test_apply_completes_when_close_races_a_deep_backlog(self):
        """With a deep backlog still draining, close() racing the applier must
        not strand it.  Termination is out-of-band (reader-done flag or closing
        flag), never via an in-band sentinel a concurrent drainer could swallow,
        so _apply always completes and its executor thread is freed."""
        loop = _loop_in_thread()
        try:
            sub = _new_subscription(_fresh_model(), loop)
            # Several iterations to shake out timing.  Each gets a fresh hand-off
            # queue + cleared closing flag (a real run() never reuses either).
            for _ in range(5):
                sub._queue = HandoffQueue(_SUBSCRIPTION_QUEUE_MAXSIZE)
                sub._closing.clear()
                model = _fresh_model()
                sub._cbk = model.handle_ref_update2
                msgs = _build_stream(2000, 8000)
                ops = _extract_ops(model, msgs)
                for op in ops:
                    sub._queue.put(op)
                applier = asyncio.run_coroutine_threadsafe(sub._apply(), loop)
                # Race a close() against the draining applier.
                sub.close()
                sub._queue.producer_done()
                # Must finish (no hang): bounded wait, not the default forever.
                applier.result(timeout=10.0)
        finally:
            loop.call_soon_threadsafe(loop.stop)

    def test_close_stops_applier_without_draining_whole_backlog(self):
        """A bare close() (no run() cancellation) must stop the applier promptly
        rather than grind through the entire queued backlog."""
        loop = _loop_in_thread()
        try:
            sub = _new_subscription(_fresh_model(), loop)
            applied = {"n": 0}

            async def _slow_cbk(*_op):
                applied["n"] += 1
                await asyncio.sleep(0)  # yield so close() can interleave

            sub._cbk = _slow_cbk
            backlog = 50000
            for i in range(backlog):
                sub._queue.put(("SET", i))
            # Reader is NOT done: only close() (closing flag) can stop _apply.
            applier = asyncio.run_coroutine_threadsafe(sub._apply(), loop)
            time.sleep(0.05)  # let it apply a bit
            sub.close()
            applier.result(timeout=10.0)
            # Stopped well before draining the whole backlog.
            self.assertLess(applied["n"], backlog)
            self.assertTrue(sub._closing.is_set())
        finally:
            loop.call_soon_threadsafe(loop.stop)


class TestDrainBatchCap(unittest.TestCase):
    """The drain must cap each batch at _MAX_DRAIN_BATCH so a deep backlog can't
    monopolize the event loop in one turn (review comment #3)."""

    def test_drain_batch_respects_max(self):
        sub = _drain_batch_subscription(maxsize=_MAX_DRAIN_BATCH * 4)
        for i in range(_MAX_DRAIN_BATCH * 3):
            sub._queue.put(("SET", i))
        batch, done = sub._drain_batch()
        self.assertFalse(done)
        self.assertEqual(len(batch), _MAX_DRAIN_BATCH)

    def test_close_during_idle_drain_tears_down_immediately(self):
        """End-to-end: with the applier parked in an IDLE drain (empty queue,
        reader not done), close() must wake it OUT OF BAND and let it terminate
        at once -- not after any poll/timeout."""
        loop = _loop_in_thread()
        try:
            sub = _new_subscription(_fresh_model(), loop)
            sub._cbk = _fresh_model().handle_ref_update2
            applier = asyncio.run_coroutine_threadsafe(sub._apply(), loop)
            time.sleep(0.3)  # ensure it's parked in drain()
            t0 = time.time()
            sub.close()  # stop() wakes the parked drain immediately
            applier.result(timeout=2.0)
            woke_in = time.time() - t0
            self.assertLess(woke_in, 0.5,
                            "idle drain not woken promptly (out-of-band) by close()")
        finally:
            loop.call_soon_threadsafe(loop.stop)


class TestApplyErrorThrottling(unittest.TestCase):
    """A recurring callback-apply failure must NOT emit a traceback per op
    (review comment #6): the traceback is logged throttled, not at
    batch-drain rate."""

    def test_repeated_failures_count_but_do_not_spam_logs(self):
        loop = _loop_in_thread()
        try:
            sub = _new_subscription(_fresh_model(), loop)

            async def _boom(*_args):
                raise RuntimeError("poison op")

            sub._cbk = _boom

            n = 5000
            for i in range(n):
                sub._queue.put(("SET", i))
            sub._queue.producer_done()

            # Capture log records emitted while applying the failing batch.
            with self.assertLogs(level="ERROR") as cm:
                asyncio.run_coroutine_threadsafe(sub._apply(),
                                                 loop).result(timeout=10.0)

            # Every failed op is skipped, but the traceback is logged a tiny,
            # throttled number of times (here a fast burst -> just the first
            # occurrence), never per-op.
            self.assertLessEqual(len(cm.records), 2)
            self.assertGreaterEqual(len(cm.records), 1)
            # The throttled log carries a traceback for debuggability.
            self.assertTrue(any(r.exc_info for r in cm.records))
        finally:
            loop.call_soon_threadsafe(loop.stop)

    def test_pending_tail_is_flushed_on_stop(self):
        """A transient burst that stops before the throttle interval elapses must
        still surface its tail count on teardown, rather than sitting unreported
        in _apply_errors_since_log until a next error that never comes."""
        loop = _loop_in_thread()
        try:
            sub = _new_subscription(_fresh_model(), loop)

            async def _boom(*_args):
                raise RuntimeError("poison op")

            sub._cbk = _boom

            n = 10
            for i in range(n):
                sub._queue.put(("SET", i))
            sub._queue.producer_done()

            with self.assertLogs(level="ERROR") as cm:
                asyncio.run_coroutine_threadsafe(sub._apply(),
                                                 loop).result(timeout=10.0)

            # The first op logs immediately (with traceback); the remaining n-1
            # accumulate within the throttle window and are flushed in one summary
            # on teardown -- nothing left pending.
            self.assertEqual(sub._apply_errors_since_log, 0)
            flush_records = [r for r in cm.records if not r.exc_info]
            self.assertTrue(flush_records, "no flush summary emitted on stop")
            self.assertTrue(any(str(n - 1) in r.getMessage() for r in flush_records),
                            "flush summary did not report the pending tail count")
        finally:
            loop.call_soon_threadsafe(loop.stop)


class TestRunDeterministicShutdown(unittest.TestCase):
    """run() must not orphan a sibling when one of reader/_apply raises.  An
    unexpected (non-CancelledError) failure is logged loudly and swallowed -- the
    subscription stops but the agent stays up -- rather than propagated to the
    CancelledError-only callers or silently dropped."""

    def _run_sub(self, loop):
        sub = _new_subscription(_fresh_model(), loop)
        # Reader does nothing but mark the queue producer_done so a normal _apply
        # terminates.
        sub._run = lambda: sub._queue.producer_done()
        return sub

    def test_apply_failure_is_logged_and_swallowed_and_stops_reader(self):
        loop = _loop_in_thread()
        try:
            sub = self._run_sub(loop)

            async def _boom_apply():
                raise RuntimeError("apply blew up")

            sub._apply = _boom_apply
            fut = asyncio.run_coroutine_threadsafe(sub.run(), loop)
            with self.assertLogs(level="ERROR") as cm:
                # run() returns normally (no propagation to the caller).
                self.assertIsNone(fut.result(timeout=10.0))
            self.assertTrue(any(r.exc_info for r in cm.records))
            # The reader was told to stop (no orphan parked on the queue).
            self.assertTrue(sub._closing.is_set())
        finally:
            loop.call_soon_threadsafe(loop.stop)

    def test_reader_failure_is_logged_and_swallowed_without_orphan(self):
        loop = _loop_in_thread()
        try:
            sub = self._run_sub(loop)

            def _boom_reader():
                # Mark done so _apply can terminate, then raise like a worker.
                sub._queue.producer_done()
                raise RuntimeError("reader blew up")

            sub._run = _boom_reader
            fut = asyncio.run_coroutine_threadsafe(sub.run(), loop)
            with self.assertLogs(level="ERROR") as cm:
                self.assertIsNone(fut.result(timeout=10.0))
            self.assertTrue(any(r.exc_info for r in cm.records))
            self.assertTrue(sub._closing.is_set())
        finally:
            loop.call_soon_threadsafe(loop.stop)

    def test_cancellederror_propagates(self):
        """A CancelledError (loop teardown / caller cancellation) is re-raised so
        cancellation propagates cleanly, unlike unexpected exceptions."""
        loop = _loop_in_thread()
        try:
            sub = self._run_sub(loop)

            async def _cancelled_apply():
                raise asyncio.CancelledError()

            sub._apply = _cancelled_apply
            fut = asyncio.run_coroutine_threadsafe(sub.run(), loop)
            # The cross-thread future surfaces cancellation as
            # concurrent.futures.CancelledError.
            with self.assertRaises(concurrent.futures.CancelledError):
                fut.result(timeout=10.0)
            self.assertTrue(sub._closing.is_set())
        finally:
            loop.call_soon_threadsafe(loop.stop)

    def test_applier_failure_unwedges_reader_blocked_in_read(self):
        """A reader parked in the blocking gRPC read must be freed by the
        iterator cancel in the shutdown path (the closing flag alone is only
        seen at the top of the _run loop).  Drive the REAL _run against a fake
        iterator whose next() blocks until cancelled, while _apply raises -->
        run() must complete and the reader thread must exit."""
        loop = _loop_in_thread()
        try:
            sub = _new_subscription(_fresh_model(), loop)
            sub._handler = None  # so _run calls self._iter.next() directly

            class _BlockingIter(object):

                def __init__(self):
                    self._cancelled = threading.Event()
                    self.parked = threading.Event()

                def next(self):
                    # Block as a real gRPC read would, until cancelled.
                    self.parked.set()
                    if self._cancelled.wait(timeout=8.0):
                        raise _CancelledRpcError()
                    raise _CancelledRpcError()

                def cancel(self):
                    self._cancelled.set()

            it = _BlockingIter()
            sub._iter = it
            reader_done = threading.Event()
            real_run = sub._run

            def _run_tracking():
                try:
                    real_run()
                finally:
                    reader_done.set()

            sub._run = _run_tracking

            async def _boom_apply():
                # Wait until the reader is genuinely parked in next(), then fail.
                while not it.parked.is_set():
                    await asyncio.sleep(0.01)
                raise RuntimeError("apply blew up while reader parked in read")

            sub._apply = _boom_apply
            fut = asyncio.run_coroutine_threadsafe(sub.run(), loop)
            with self.assertLogs(level="ERROR"):
                # run() logs-and-swallows the unexpected error; returns normally.
                self.assertIsNone(fut.result(timeout=10.0))
            # The blocking reader was unwedged by the iterator cancel.
            self.assertTrue(it._cancelled.is_set())
            self.assertTrue(reader_done.wait(timeout=3.0))
        finally:
            loop.call_soon_threadsafe(loop.stop)


class TestRunSubscriptionLifecycle(unittest.TestCase):
    """The single _run_subscription helper guarantees teardown + deregistration
    on ANY termination, so a dead subscription is never left registered."""

    def _client(self):
        from preseem.grpc_client import PreseemGrpcClient
        return PreseemGrpcClient.__new__(PreseemGrpcClient)

    def _fake_sub(self, result):
        sub = Subscription.__new__(Subscription)
        sub.closed = False

        async def _run():
            if isinstance(result, BaseException):
                raise result
            return result

        sub.run = _run
        sub.close = lambda: setattr(sub, "closed", True)
        return sub

    def _drive(self, coro):
        loop = _loop_in_thread()
        try:
            return asyncio.run_coroutine_threadsafe(coro, loop).result(timeout=5)
        finally:
            loop.call_soon_threadsafe(loop.stop)

    def test_normal_return_deregisters_and_closes(self):
        client = self._client()
        reg = {}
        sub = self._fake_sub(None)
        reg["k"] = sub
        self._drive(client._run_subscription(sub, registry=reg, key="k"))
        self.assertTrue(sub.closed)
        self.assertNotIn("k", reg)  # deregistered on normal return (#4)

    def test_swallowed_error_return_deregisters(self):
        # run() swallows unexpected errors and returns None; helper still cleans.
        client = self._client()
        reg = {}
        sub = self._fake_sub(None)
        reg["k"] = sub
        self._drive(client._run_subscription(sub, registry=reg, key="k"))
        self.assertNotIn("k", reg)

    def test_cancellation_deregisters_and_is_suppressed(self):
        client = self._client()
        reg = {}
        sub = self._fake_sub(asyncio.CancelledError())
        reg["k"] = sub
        # Suppressed (no raise) and deregistered.
        self._drive(client._run_subscription(sub, registry=reg, key="k"))
        self.assertTrue(sub.closed)
        self.assertNotIn("k", reg)

    def test_only_removes_its_own_registry_entry(self):
        client = self._client()
        reg = {}
        old = self._fake_sub(None)
        newer = self._fake_sub(None)
        reg["k"] = newer  # a newer sub already replaced old under the same key
        self._drive(client._run_subscription(old, registry=reg, key="k"))
        self.assertIs(reg.get("k"), newer)  # newer not clobbered

    def test_no_registry_still_closes(self):
        client = self._client()
        sub = self._fake_sub(None)
        self._drive(client._run_subscription(sub))
        self.assertTrue(sub.closed)

    def test_slot_attribute_cleared_on_termination(self):
        # The nm_ent/nm_ref singletons live in a client attribute, not a dict;
        # _run_subscription must clear that slot on termination too.
        client = self._client()
        sub = self._fake_sub(None)
        client._nm_ent_subscription = sub
        self._drive(client._run_subscription(sub, slot="_nm_ent_subscription"))
        self.assertTrue(sub.closed)
        self.assertIsNone(client._nm_ent_subscription)

    def test_slot_only_cleared_if_still_this_sub(self):
        # A newer subscription replaced the slot -> the old run's teardown must
        # not null out the newer one.
        client = self._client()
        old = self._fake_sub(None)
        newer = self._fake_sub(None)
        client._nm_ent_subscription = newer
        self._drive(client._run_subscription(old, slot="_nm_ent_subscription"))
        self.assertIs(client._nm_ent_subscription, newer)


if __name__ == "__main__":
    unittest.main()
