teng-lin--notebooklm-py/tests/unit/test_logging_correlation.py

"""Tests for per-request correlation IDs via contextvars.ContextVar."""

from __future__ import annotations

import asyncio
import io
import logging
import re
from contextvars import Token

import pytest

from notebooklm._logging import (
    RedactingFilter,
    RedactingFormatter,
    get_request_id,
    reset_request_id,
    set_request_id,
)
from tests._helpers.client_factory import build_client_shell_for_tests


@pytest.fixture(autouse=True)
def clear_request_id():
    """Snapshot/restore the request id ContextVar around each test."""
    from notebooklm._logging import _current_request_id

    token = _current_request_id.set(_current_request_id.get())
    try:
        yield
    finally:
        _current_request_id.reset(token)


def _make_record(msg: str = "test") -> logging.LogRecord:
    return logging.LogRecord(
        name="notebooklm.test",
        level=logging.WARNING,
        pathname=__file__,
        lineno=0,
        msg=msg,
        args=None,
        exc_info=None,
    )


# ---------------------------------------------------------------------------
# set/reset/get primitives
# ---------------------------------------------------------------------------


def test_set_request_id_returns_token_and_generated_id():
    """set_request_id() returns a Token; get_request_id() returns 8-hex."""
    token = set_request_id()
    try:
        assert isinstance(token, Token)
        rid = get_request_id()
        assert rid is not None
        assert re.fullmatch(r"[0-9a-f]{8}", rid)
    finally:
        reset_request_id(token)


def test_set_request_id_accepts_explicit_value():
    """Pass req_id=... to set a specific id."""
    token = set_request_id("CUSTOM01")
    try:
        assert get_request_id() == "CUSTOM01"
    finally:
        reset_request_id(token)


def test_reset_restores_previous_value():
    """reset_request_id restores to PREVIOUS, not None."""
    outer = set_request_id("OUTER001")
    try:
        inner = set_request_id("INNER001")
        try:
            assert get_request_id() == "INNER001"
        finally:
            reset_request_id(inner)
        assert get_request_id() == "OUTER001"
    finally:
        reset_request_id(outer)
    assert get_request_id() is None


# ---------------------------------------------------------------------------
# Filter integration
# ---------------------------------------------------------------------------


def test_filter_prefixes_after_scrub():
    """Filter prefixes [req=<id>] AFTER scrubbing. Token never appears in output."""
    filt = RedactingFilter()
    token = set_request_id("ABCD0123")
    try:
        rec = _make_record("scrubbable at=SECRET_TOK detail")
        filt.filter(rec)
        # Prefix outside the scrubbed body
        assert rec.msg.startswith("[req=ABCD0123] ")
        # Scrub still happened
        assert "SECRET_TOK" not in rec.msg
        assert "at=***" in rec.msg
    finally:
        reset_request_id(token)


def test_filter_no_prefix_when_id_unset():
    """With no id set, record.msg has no prefix."""
    filt = RedactingFilter()
    rec = _make_record("plain message")
    filt.filter(rec)
    assert not rec.msg.startswith("[req=")
    assert rec.msg == "plain message"


def test_filter_no_double_prefix_through_two_filters():
    """Marker attribute prevents second filter pass from re-prefixing."""
    filt1 = RedactingFilter()
    filt2 = RedactingFilter()
    token = set_request_id("XY00ZZ11")
    try:
        rec = _make_record("once")
        filt1.filter(rec)
        filt2.filter(rec)  # second pass, simulating two handlers
        # Only one [req=...] prefix
        assert rec.msg.count("[req=XY00ZZ11]") == 1
    finally:
        reset_request_id(token)


def test_filter_includes_prefix_for_exc_info_records():
    """exc_info path also gets the prefix on the message line."""
    filt = RedactingFilter()
    token = set_request_id("EXC00001")
    try:
        try:
            raise ValueError("oops at=LEAK")
        except ValueError:
            import sys

            rec = logging.LogRecord(
                name="notebooklm.test",
                level=logging.ERROR,
                pathname=__file__,
                lineno=0,
                msg="failure",
                args=None,
                exc_info=sys.exc_info(),
            )
        filt.filter(rec)
        assert rec.msg.startswith("[req=EXC00001] ")
        assert "LEAK" not in (rec.exc_text or "")
    finally:
        reset_request_id(token)


# ---------------------------------------------------------------------------
# asyncio Task isolation
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_concurrent_tasks_see_distinct_ids():
    """gather(task_a, task_b) — each sets its own id; emissions are tagged distinctly."""
    captured: dict[str, str] = {}

    async def worker(name: str, explicit_id: str) -> None:
        token = set_request_id(explicit_id)
        try:
            await asyncio.sleep(0)  # yield to interleave with sibling
            captured[name] = get_request_id() or ""
        finally:
            reset_request_id(token)

    await asyncio.gather(worker("a", "AAAA0000"), worker("b", "BBBB1111"))
    assert captured["a"] == "AAAA0000"
    assert captured["b"] == "BBBB1111"


@pytest.mark.asyncio
async def test_id_does_not_leak_after_reset():
    """After reset, subsequent emissions have no prefix."""
    filt = RedactingFilter()
    token = set_request_id("KEEP0001")
    rec1 = _make_record("inside")
    filt.filter(rec1)
    assert "[req=KEEP0001]" in rec1.msg

    reset_request_id(token)
    rec2 = _make_record("outside")
    filt.filter(rec2)
    assert "[req=" not in rec2.msg


@pytest.mark.asyncio
async def test_correlation_threads_through_child_loggers():
    """Emissions from notebooklm._core and notebooklm._chat get same prefix
    within a single task."""
    filt = RedactingFilter()
    token = set_request_id("CHILD001")
    try:
        rec_core = logging.LogRecord(
            name="notebooklm._core",
            level=logging.WARNING,
            pathname=__file__,
            lineno=0,
            msg="core msg",
            args=None,
            exc_info=None,
        )
        rec_chat = logging.LogRecord(
            name="notebooklm._chat",
            level=logging.WARNING,
            pathname=__file__,
            lineno=0,
            msg="chat msg",
            args=None,
            exc_info=None,
        )
        filt.filter(rec_core)
        filt.filter(rec_chat)
        assert rec_core.msg.startswith("[req=CHILD001] ")
        assert rec_chat.msg.startswith("[req=CHILD001] ")
    finally:
        reset_request_id(token)


# ---------------------------------------------------------------------------
# Integration with full Formatter + Filter
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_retry_inherits_parent_request_id():
    """Recursive executor.rpc_call(_is_retry=True) must NOT mint a fresh id — the
    failure→refresh→retry sequence should appear under one prefix."""
    from notebooklm.auth import AuthTokens

    captured_ids: list[str | None] = []

    async def fake_impl(
        method,
        params,
        source_path,
        allow_null,
        is_retry,
        *,
        disable_internal_retries: bool = False,
        operation_variant: str | None = None,
        _refresh_budget=None,
        _retry_deadline=None,
    ):
        captured_ids.append(get_request_id())
        # First call: raise to trigger retry path; second call: succeed.
        if not is_retry:
            # Mimic decode-time retry without leaving the executor.
            return await executor.rpc_call(
                method,
                params,
                source_path,
                allow_null,
                _is_retry=True,
            )
        return "ok"

    # Real NotebookLMClient shell — the executor's open-client guard
    # requires a truthy http_client, so we ``open()`` and let the lifecycle
    # construct one against the default httpx transport. ``fake_impl`` is
    # monkeypatched onto ``_execute_once`` so no actual HTTP call fires;
    # the test exercises the request-id propagation through the executor
    # wrapper purely in-process.
    auth = AuthTokens(cookies={"SID": "test_sid"}, csrf_token="csrf", session_id="sid")
    core = build_client_shell_for_tests(auth)
    await core.__aenter__()
    try:
        executor = core._rpc_executor
        executor._execute_once = fake_impl  # type: ignore[method-assign]

        result = await core._rpc_executor.rpc_call(method=object(), params=[])  # type: ignore[arg-type]
        assert result == "ok"
        assert len(captured_ids) == 2
        assert captured_ids[0] == captured_ids[1]
        assert captured_ids[0] is not None
    finally:
        await core.close()


def test_end_to_end_prefix_visible_in_rendered_output():
    """When wired to a real handler, the rendered output has [req=...] prefix."""
    buf = io.StringIO()
    handler = logging.StreamHandler(buf)
    handler.setLevel(logging.NOTSET)
    handler.setFormatter(RedactingFormatter(logging.Formatter("%(message)s")))
    handler.addFilter(RedactingFilter())

    logger = logging.getLogger("test_e2e_corr_unique")
    logger.handlers.clear()
    logger.filters.clear()
    logger.addHandler(handler)
    logger.setLevel(logging.WARNING)
    logger.propagate = False

    token = set_request_id("E2EE0001")
    try:
        logger.warning("hello with at=SECRET")
    finally:
        reset_request_id(token)

    out = buf.getvalue()
    assert "[req=E2EE0001]" in out
    assert "SECRET" not in out
    assert "at=***" in out

    logger.handlers.clear()
    logger.filters.clear()