news/tests/test_cninfo.py

"""cninfo 抓取并发限流测试。

背景:
    每次 `_render_page` 都会启动一个完整的 headless Chromium(实测约 850MB /
    10 进程),而 `crawl_watchlist` 对 15 只股票全量并发 → 峰值需求 ≈12.5GB,
    远超本机 7.9GB RAM,曾导致 2026-09-06/08/10 三次整机冻结。
    修复方式:模块级信号量限制同时存活的浏览器数(默认 2)。

本测试用假 AsyncWebCrawler 验证并发上限,不启动真实浏览器。
"""

from __future__ import annotations

import asyncio
from types import SimpleNamespace

import pytest


def _install_fake_crawler(monkeypatch: pytest.MonkeyPatch, tracker: dict) -> None:
    """把 crawl4ai.AsyncWebCrawler 换成记录并发峰值的假实现。"""
    import crawl4ai

    class FakeCrawler:
        def __init__(self, config=None) -> None:  # noqa: ARG002
            pass

        async def __aenter__(self) -> FakeCrawler:
            tracker["current"] += 1
            tracker["max"] = max(tracker["max"], tracker["current"])
            return self

        async def __aexit__(self, *exc: object) -> bool:
            tracker["current"] -= 1
            return False

        async def arun(self, url: str, config=None) -> SimpleNamespace:  # noqa: ARG002
            # 模拟渲染耗时,制造并发窗口
            await asyncio.sleep(0.05)
            return SimpleNamespace(html=f"<html>{url}</html>")

    monkeypatch.setattr(crawl4ai, "AsyncWebCrawler", FakeCrawler)


# --------------------------------------------------------------------------- #
# 并发上限解析
# --------------------------------------------------------------------------- #

def test_resolve_concurrency_default(monkeypatch: pytest.MonkeyPatch) -> None:
    """未设 env 时使用默认值 2。"""
    monkeypatch.delenv("CNINFO_RENDER_CONCURRENCY", raising=False)
    from crawler.cninfo import DEFAULT_RENDER_CONCURRENCY, _resolve_render_concurrency

    assert DEFAULT_RENDER_CONCURRENCY == 2
    assert _resolve_render_concurrency() == 2


def test_resolve_concurrency_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
    """env CNINFO_RENDER_CONCURRENCY 可覆盖。"""
    monkeypatch.setenv("CNINFO_RENDER_CONCURRENCY", "4")
    from crawler.cninfo import _resolve_render_concurrency

    assert _resolve_render_concurrency() == 4


def test_resolve_concurrency_invalid_falls_back(monkeypatch: pytest.MonkeyPatch) -> None:
    """非法值回退默认值,不抛异常。"""
    monkeypatch.setenv("CNINFO_RENDER_CONCURRENCY", "abc")
    from crawler.cninfo import _resolve_render_concurrency

    assert _resolve_render_concurrency() == 2


def test_resolve_concurrency_clamped_to_one(monkeypatch: pytest.MonkeyPatch) -> None:
    """0/负数夹到 1,避免信号量死锁。"""
    monkeypatch.setenv("CNINFO_RENDER_CONCURRENCY", "0")
    from crawler.cninfo import _resolve_render_concurrency

    assert _resolve_render_concurrency() == 1


# --------------------------------------------------------------------------- #
# 渲染并发限流
# --------------------------------------------------------------------------- #

def test_render_concurrency_capped(monkeypatch: pytest.MonkeyPatch) -> None:
    """6 个并发渲染请求,同时存活的浏览器数不超过信号量上限 2。"""
    from crawler import cninfo

    tracker = {"current": 0, "max": 0}
    _install_fake_crawler(monkeypatch, tracker)
    # 独立信号量,避免跨测试污染模块级状态
    monkeypatch.setattr(cninfo, "_render_sem", asyncio.Semaphore(2))

    async def run() -> list[str]:
        return await asyncio.gather(*[
            cninfo._render_page(f"https://example.com/{i}") for i in range(6)
        ])

    results = asyncio.run(run())
    assert len(results) == 6, "全部请求都应完成(限流不应丢请求)"
    assert tracker["max"] <= 2, f"并发峰值 {tracker['max']} 超过上限 2"


def test_render_concurrency_one_is_serial(monkeypatch: pytest.MonkeyPatch) -> None:
    """上限为 1 时完全串行执行。"""
    from crawler import cninfo

    tracker = {"current": 0, "max": 0}
    _install_fake_crawler(monkeypatch, tracker)
    monkeypatch.setattr(cninfo, "_render_sem", asyncio.Semaphore(1))

    async def run() -> list[str]:
        return await asyncio.gather(*[
            cninfo._render_page(f"https://example.com/{i}") for i in range(4)
        ])

    results = asyncio.run(run())
    assert len(results) == 4
    assert tracker["max"] == 1, "上限 1 时不应出现并发"


def test_render_page_returns_html(monkeypatch: pytest.MonkeyPatch) -> None:
    """渲染正常返回 html 字符串。"""
    from crawler import cninfo

    tracker = {"current": 0, "max": 0}
    _install_fake_crawler(monkeypatch, tracker)
    monkeypatch.setattr(cninfo, "_render_sem", asyncio.Semaphore(2))

    html = asyncio.run(cninfo._render_page("https://example.com/x"))
    assert html == "<html>https://example.com/x</html>"


def test_render_semaphore_released_on_error(monkeypatch: pytest.MonkeyPatch) -> None:
    """渲染抛异常时也要释放槽位(否则后续渲染会死锁)。"""
    import crawl4ai

    from crawler import cninfo

    class BoomCrawler:
        def __init__(self, config=None) -> None:  # noqa: ARG002
            pass

        async def __aenter__(self) -> BoomCrawler:
            raise RuntimeError("浏览器启动失败")

        async def __aexit__(self, *exc: object) -> bool:
            return False

        async def arun(self, url: str, config=None):  # noqa: ARG002
            raise RuntimeError("unreachable")

    monkeypatch.setattr(crawl4ai, "AsyncWebCrawler", BoomCrawler)
    monkeypatch.setattr(cninfo, "_render_sem", asyncio.Semaphore(1))

    async def run() -> None:
        # 第一次失败后,槽位必须已释放,第二次才能拿到
        with pytest.raises(RuntimeError):
            await cninfo._render_page("https://example.com/fail")
        with pytest.raises(RuntimeError):
            await asyncio.wait_for(
                cninfo._render_page("https://example.com/fail2"), timeout=2
            )

    asyncio.run(run())
    assert cninfo._render_sem._value == 1, "异常后槽位未释放"