content-find-agent-new/tests/test_unified_search_unit.py

"""M8E:统一搜索单元(游走复用 _ProgressiveContext)——前3命中再翻页、30s 限速、seen/index 注入共享。"""

from typing import Any

from content_agent.business_modules import progressive_screening
from content_agent.business_modules.progressive_screening import _ProgressiveContext
from tests.gemini_helpers import FakeGeminiVideoClient
from tests.p6_walk_helpers import build_initial_walk_context
from tests.test_progressive_screening import FakeClock


def _result(cid: str, *, has_more: bool = False, next_cursor: str = "") -> dict[str, Any]:
    return {
        "content_discovery_id": f"cd_{cid}",
        "search_query_id": "tagq_1",
        "platform": "douyin",
        "platform_content_id": cid,
        "platform_content_format": "video",
        "description": "x",
        "platform_author_id": f"author_{cid}",
        "author_display_name": "n",
        "statistics": {
            # M11 re-baseline:原 digg=5M/极低占比新算只得 45.08,总分 62.54<70 不入池,
            # "前3命中再翻页"不触发。改为高共鸣占比 → 平台分 ~88、总分 ~84 真入池,游走分页生效。
            "digg_count": 100_000,
            "comment_count": 8_000,
            "share_count": 30_000,
            "collect_count": 50_000,
        },
        "tags": [],
        "score": 72,
        "risk_level": "low",
        "availability": "available",
        "discovery_relation": "fake",
        "discovery_start_source": "pattern_itemset",
        "previous_discovery_step": "hashtag_to_query",
        "content_metadata_source": "fake",
        "platform_raw_payload": {"content_id": cid},
        "has_more": has_more,
        "next_cursor": next_cursor,
    }


class _PagedClient:
    def __init__(self, pages: dict[str, list[dict[str, Any]]]) -> None:
        self.pages = pages
        self.calls: list[dict[str, Any]] = []

    def search_full_page(self, query: dict[str, Any]) -> list[dict[str, Any]]:
        self.calls.append(dict(query))
        return list(self.pages.get(str(query.get("page_cursor") or ""), []))

    def search(self, query: dict[str, Any]) -> list[dict[str, Any]]:
        return self.search_full_page(query)


class _OnePerCallClient:
    def __init__(self) -> None:
        self.calls: list[dict[str, Any]] = []
        self.n = 0

    def search(self, query: dict[str, Any]) -> list[dict[str, Any]]:
        self.calls.append(dict(query))
        self.n += 1
        return [_result(f"rc{self.n}")]


def _tag_query(query_id: str) -> dict[str, Any]:
    return {
        "record_schema_version": "runtime_record.v1",
        "run_id": "run_001",
        "policy_run_id": "policy_run_001",
        "search_query_id": query_id,
        "search_query": "人物故事",
        "search_query_generation_method": "tag_query",
        "discovery_start_source": "pattern_itemset",
        "previous_discovery_step": "hashtag_to_query",
    }


def _walk_ctx(base, *, client, limiter=None, seen=None, recall_base=0, decision_base=0):
    return _ProgressiveContext(
        run_id=base["run_id"],
        policy_run_id=base["policy_run_id"],
        source_context=base["source_context"],
        policy_bundle=base["policy_bundle"],
        platform_client=client,
        runtime=base["runtime"],
        gemini_video_client=FakeGeminiVideoClient(),
        limiter=limiter,
        archive_dispatcher=None,
        external_seen_content_ids=seen if seen is not None else set(),
        recall_index_base=recall_base,
        decision_index_base=decision_base,
    )


def test_walk_search_unit_is_progressive(tmp_path):
    base = build_initial_walk_context(tmp_path)
    client = _PagedClient(
        {
            "": [
                _result("c1", has_more=True, next_cursor="cur2"),
                _result("c2"),
                _result("c3"),
                _result("c4"),
                _result("c5"),
            ],
            "cur2": [_result("c6"), _result("c7")],
        }
    )
    walk_ctx = _walk_ctx(base, client=client)

    batch = walk_ctx.process_query(_tag_query("tagq_1"))

    # 前3命中(c1/c2/c3 入池)→ 处理剩余(c4/c5)→ 翻页拉 cur2(c6/c7)。
    assert [str(c.get("page_cursor") or "") for c in client.calls] == ["", "cur2"]
    ids = [d["decision_target_id"] for d in batch["rule_decisions"]]
    assert ids == ["c1", "c2", "c3", "c4", "c5", "c6", "c7"]


def test_walk_search_unit_rate_limited(tmp_path):
    base = build_initial_walk_context(tmp_path)
    clock = FakeClock()
    limiter = progressive_screening.SearchCallLimiter(
        now_fn=clock.monotonic, sleep_fn=clock.sleep
    )
    walk_ctx = _walk_ctx(base, client=_OnePerCallClient(), limiter=limiter)

    walk_ctx.process_query(_tag_query("tagq_a"))
    before = clock.now
    walk_ctx.process_query(_tag_query("tagq_b"))

    # 第二次搜索受限速器约束:相邻搜索间隔 ≥30s(fake clock 推进 ≥30)。
    assert clock.now - before >= 30.0


def test_walk_search_unit_shares_seen_and_index(tmp_path):
    base = build_initial_walk_context(tmp_path)
    seen = {"cX"}  # 注入"首轮已发现"
    client = _PagedClient({"": [_result("cX"), _result("cY")]})
    walk_ctx = _walk_ctx(base, client=client, seen=seen, recall_base=20, decision_base=20)

    batch = walk_ctx.process_query(_tag_query("tagq_1"))

    # 注入 seen 跨调用去重:cX 跳过,只判 cY;cY 回写进共享 seen 集。
    assert [d["decision_target_id"] for d in batch["rule_decisions"]] == ["cY"]
    assert "cY" in seen
    # index 基线生效:游走 id 从基线起算,不撞首轮(decision_base=20 → 首条 d_021)。
    assert batch["rule_decisions"][0]["decision_id"] == "d_021"