simon
7b33182f67
- P1-2: 提取 DEFAULT_NEWS_STEPS(全链路去 report+cninfo),定时/补跑/默认三处统一; 启动补跑不再误执行 cninfo 公告管道 - P1-3: 新增 scheduler/timeutil.py(调度时区统一入口,SCHEDULE_TZ 可覆盖, 默认 Asia/Shanghai);run_scheduler 定时/补跑/--once、crawler 补跑保护、 reporter 兜底日期全部改用调度时区,避免系统时区非上海时日期错位 - 新增 6 个测试;全量 282 passed
452 lines
17 KiB
Python
452 lines
17 KiB
Python
"""增量处理与 pipeline 断点续跑测试。
|
|
|
|
覆盖:
|
|
- M2/M4/M5 脚本的「产物存在即跳过」过滤逻辑
|
|
- scheduler.pipeline 断点状态记录与 --resume 续跑逻辑
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import UTC
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# M4 / M5: 产物存在即跳过
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_llm_filter_existing_skips_done(tmp_path: Path) -> None:
|
|
"""M4:输出目录已有 {url_hash}.json 的输入被过滤,不重复调用 LLM API。"""
|
|
from scripts.run_event_extraction import _filter_existing
|
|
|
|
out_dir = tmp_path / "out"
|
|
out_dir.mkdir()
|
|
# 已处理
|
|
(out_dir / "aaa.json").write_text("{}", encoding="utf-8")
|
|
(out_dir / "ccc.json").write_text("{}", encoding="utf-8")
|
|
files = [
|
|
tmp_path / "in" / "aaa.json", # 已处理 → 跳过
|
|
tmp_path / "in" / "bbb.json", # 未处理 → 待处理
|
|
tmp_path / "in" / "ccc.json", # 已处理 → 跳过
|
|
]
|
|
pending, skipped = _filter_existing(files, out_dir)
|
|
assert skipped == 2
|
|
assert [p.stem for p in pending] == ["bbb"]
|
|
|
|
|
|
def test_llm_filter_existing_force_keeps_all(tmp_path: Path) -> None:
|
|
"""M4:--force 时不做过滤(全量重抽由调用方控制)。"""
|
|
from scripts.run_event_extraction import _filter_existing
|
|
|
|
out_dir = tmp_path / "out"
|
|
out_dir.mkdir()
|
|
(out_dir / "aaa.json").write_text("{}", encoding="utf-8")
|
|
files = [tmp_path / "in" / "aaa.json"]
|
|
# _filter_existing 本身不含 force 逻辑,验证在 force 下不会被调用:
|
|
# 直接验证「已存在也被返回」需由上层跳过调用,这里仅确认过滤函数行为。
|
|
pending, skipped = _filter_existing(files, out_dir)
|
|
assert skipped == 1
|
|
assert pending == []
|
|
|
|
|
|
def test_embedding_filter_existing_skips_done(tmp_path: Path) -> None:
|
|
"""M5:输出目录已有 {url_hash}.json 的输入被过滤,不重复调用 embed API。"""
|
|
from scripts.run_embedding import _filter_existing
|
|
|
|
out_dir = tmp_path / "emb"
|
|
out_dir.mkdir()
|
|
(out_dir / "h1.json").write_text("{}", encoding="utf-8")
|
|
files = [
|
|
(tmp_path / "in" / "h1.json", "event"), # 已处理 → 跳过
|
|
(tmp_path / "in" / "h2.json", "event"), # 未处理 → 待处理
|
|
]
|
|
pending, skipped = _filter_existing(files, out_dir)
|
|
assert skipped == 1
|
|
assert [p.stem for p, _ in pending] == ["h2"]
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# M2: 已提取文章跳过
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_extractor_process_source_day_skips_existing(tmp_path: Path) -> None:
|
|
"""M2:输出目录已有产物的记录被跳过提取,且 index 回补完整。"""
|
|
from scripts.run_extractor import _process_source_day
|
|
|
|
raw_dir = tmp_path / "raw" / "cls" / "20260616"
|
|
raw_dir.mkdir(parents=True)
|
|
# 两条 raw 记录(url_hash 与产物文件名一致)
|
|
rec1 = {"source_id": "cls", "url": "https://a/1", "url_hash": "aaa1111111111111",
|
|
"stage": "article", "success": True, "html_file": "aaa1111111111111.html"}
|
|
rec2 = {"source_id": "cls", "url": "https://b/2", "url_hash": "bbb2222222222222",
|
|
"stage": "article", "success": True, "html_file": "bbb2222222222222.html"}
|
|
with (raw_dir / "index.jsonl").open("a", encoding="utf-8") as f:
|
|
f.write(json.dumps(rec1, ensure_ascii=False) + "\n")
|
|
f.write(json.dumps(rec2, ensure_ascii=False) + "\n")
|
|
|
|
out_dir = tmp_path / "proc" / "cls" / "20260616"
|
|
out_dir.mkdir(parents=True)
|
|
# 预置一条已有产物(视为已提取)
|
|
article = {
|
|
"source_id": "cls", "url": "https://a/1", "url_hash": "aaa1111111111111",
|
|
"title": "已有", "content": "内容", "word_count": 2,
|
|
}
|
|
(out_dir / "aaa1111111111111.json").write_text(
|
|
json.dumps(article, ensure_ascii=False), encoding="utf-8"
|
|
)
|
|
|
|
# 另一条无 html 文件 → 提取失败(但不影响跳过逻辑断言)
|
|
succ, total, skipped = _process_source_day(
|
|
"cls", "20260616", tmp_path / "raw", tmp_path / "proc"
|
|
)
|
|
assert total == 2
|
|
assert skipped == 1 # 已有产物被跳过
|
|
assert succ == 0 # 另一条因 html 缺失提取失败
|
|
# index 回补了被跳过条目的行
|
|
idx = (out_dir / "index.jsonl").read_text(encoding="utf-8").strip()
|
|
assert "aaa1111111111111" in idx
|
|
|
|
|
|
def test_extractor_process_source_day_force_rebuilds(tmp_path: Path) -> None:
|
|
"""M2:--force 时不做跳过,并重建 index。"""
|
|
from scripts.run_extractor import _process_source_day
|
|
|
|
raw_dir = tmp_path / "raw" / "cls" / "20260616"
|
|
raw_dir.mkdir(parents=True)
|
|
rec = {"source_id": "cls", "url": "https://a/1", "url_hash": "aaa1111111111111",
|
|
"stage": "article", "success": True, "html_file": "aaa1111111111111.html"}
|
|
with (raw_dir / "index.jsonl").open("a", encoding="utf-8") as f:
|
|
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
|
|
|
out_dir = tmp_path / "proc" / "cls" / "20260616"
|
|
out_dir.mkdir(parents=True)
|
|
(out_dir / "aaa1111111111111.json").write_text("{}", encoding="utf-8")
|
|
(out_dir / "index.jsonl").write_text("旧内容", encoding="utf-8")
|
|
|
|
succ, total, skipped = _process_source_day(
|
|
"cls", "20260616", tmp_path / "raw", tmp_path / "proc", force=True
|
|
)
|
|
assert skipped == 0
|
|
# force 模式重建 index(旧内容被清掉;此处无 html 提取失败,index 为空或不存在)
|
|
idx = out_dir / "index.jsonl"
|
|
assert not idx.exists() or idx.read_text(encoding="utf-8") == ""
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# pipeline 断点状态与 --resume
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@pytest.fixture
|
|
def fake_subprocess(monkeypatch: pytest.MonkeyPatch):
|
|
"""mock subprocess.run,按步骤名返回 returncode,并记录调用顺序。"""
|
|
from scheduler import pipeline
|
|
|
|
calls: list[str] = []
|
|
|
|
def _step_name(cmd: list[str]) -> str:
|
|
"""从命令中提取脚本名,如 scripts.run_extractor → run_extractor。"""
|
|
return next(c.split(".")[-1] for c in cmd if "scripts.run_" in c)
|
|
|
|
def _fake_run(cmd, timeout=None): # noqa: ARG001
|
|
calls.append(_step_name(cmd))
|
|
return SimpleNamespace(returncode=0)
|
|
|
|
monkeypatch.setattr(pipeline.subprocess, "run", _fake_run)
|
|
return calls
|
|
|
|
|
|
def _run_with_steps(monkeypatch: pytest.MonkeyPatch, failures: set[str]):
|
|
"""构造 run_step:指定步骤(如 'extractor')返回失败。"""
|
|
from scheduler import pipeline
|
|
|
|
def _fake_run(cmd, timeout=None): # noqa: ARG001
|
|
name = next(c.split(".")[-1] for c in cmd if "scripts.run_" in c)
|
|
name = name.replace("run_", "") # run_extractor → extractor
|
|
return SimpleNamespace(returncode=1 if name in failures else 0)
|
|
|
|
monkeypatch.setattr(pipeline.subprocess, "run", _fake_run)
|
|
|
|
|
|
def test_pipeline_records_state(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""全量运行后状态文件按日期记录每个步骤的 ok/failed。"""
|
|
from scheduler import pipeline
|
|
|
|
_run_with_steps(monkeypatch, failures={"extractor"})
|
|
state_path = tmp_path / "state.json"
|
|
steps = ["extractor", "dedup", "llm"]
|
|
pipeline.run_pipeline("20260616", steps=steps, state_path=state_path)
|
|
|
|
state = pipeline._load_pipeline_state(state_path)
|
|
day = state["20260616"]
|
|
assert day["extractor"]["status"] == "failed"
|
|
assert day["dedup"]["status"] == "ok"
|
|
assert day["llm"]["status"] == "ok"
|
|
# dedup 返回 1 被特判为成功,故用 extractor 制造失败
|
|
assert day["extractor"]["exit_code"] == 1
|
|
|
|
|
|
def test_pipeline_resume_skips_success_prefix(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""resume 跳过连续成功步骤,从失败步骤继续。"""
|
|
from scheduler import pipeline
|
|
|
|
state_path = tmp_path / "state.json"
|
|
# 预置状态:extractor 失败,dedup/llm 成功(模拟上次运行)
|
|
state = {"20260616": {
|
|
"extractor": {"status": "failed", "exit_code": 1},
|
|
"dedup": {"status": "ok", "exit_code": 0},
|
|
"llm": {"status": "ok", "exit_code": 0},
|
|
}}
|
|
pipeline._save_pipeline_state(state, state_path)
|
|
|
|
calls: list[str] = []
|
|
|
|
def _fake_run(cmd, timeout=None): # noqa: ARG001
|
|
name = next(c.split(".")[-1] for c in cmd if "scripts.run_" in c)
|
|
calls.append(name)
|
|
return SimpleNamespace(returncode=0)
|
|
|
|
monkeypatch.setattr(pipeline.subprocess, "run", _fake_run)
|
|
steps = ["extractor", "dedup", "llm"]
|
|
result = pipeline.run_pipeline(
|
|
"20260616", steps=steps, resume=True, state_path=state_path
|
|
)
|
|
# 从 extractor 开始重跑全部(extractor 之后的 dedup/llm 需重跑以覆盖降级数据)
|
|
assert calls == ["run_extractor", "run_dedup", "run_event_extraction"]
|
|
assert all(s.success for s in result.steps)
|
|
|
|
|
|
def test_pipeline_resume_all_done_noop(tmp_path: Path) -> None:
|
|
"""resume 且所有步骤均已成功时,不执行任何步骤。"""
|
|
from scheduler import pipeline
|
|
|
|
state_path = tmp_path / "state.json"
|
|
state = {"20260616": {
|
|
"extractor": {"status": "ok", "exit_code": 0},
|
|
"dedup": {"status": "ok", "exit_code": 0},
|
|
"llm": {"status": "ok", "exit_code": 0},
|
|
}}
|
|
pipeline._save_pipeline_state(state, state_path)
|
|
|
|
steps = ["extractor", "dedup", "llm"]
|
|
result = pipeline.run_pipeline(
|
|
"20260616", steps=steps, resume=True, state_path=state_path
|
|
)
|
|
assert result.steps == []
|
|
assert result.all_success # 空步骤视为成功
|
|
|
|
|
|
def test_pipeline_resume_missing_step_starts_from_first_missing(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""resume:部分步骤无历史记录时,从首个缺失步骤开始。"""
|
|
from scheduler import pipeline
|
|
|
|
state_path = tmp_path / "state.json"
|
|
state = {"20260616": {
|
|
"extractor": {"status": "ok", "exit_code": 0},
|
|
}}
|
|
pipeline._save_pipeline_state(state, state_path)
|
|
|
|
idx = pipeline._resume_start_index(
|
|
["extractor", "dedup", "llm"], "20260616",
|
|
pipeline._load_pipeline_state(state_path),
|
|
)
|
|
assert idx == 1 # dedup 缺失 → 从它开始
|
|
|
|
|
|
def test_once_rejects_resume_with_steps(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""--resume 与 --steps 同时使用时报错。"""
|
|
import scripts.run_scheduler as rs
|
|
|
|
args = SimpleNamespace(steps="crawler,extractor", resume=True,
|
|
date="20260616")
|
|
rc = rs._once(args)
|
|
assert rc == 2
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# P0-3: 补跑模式跳过抓取 / P0-4: report 使用传入日期
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_crawler_backfill_skips_past_date(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""补跑历史日期时跳过 crawler(首页只含当天内容),不执行子进程。"""
|
|
from scheduler import pipeline
|
|
|
|
def _fail_run(cmd, timeout=None): # noqa: ARG001
|
|
raise AssertionError("补跑模式不应执行子进程")
|
|
|
|
monkeypatch.setattr(pipeline.subprocess, "run", _fail_run)
|
|
sr = pipeline.run_step("crawler", "20260101")
|
|
assert sr.success is True
|
|
assert "跳过" in sr.tail_msg
|
|
assert sr.elapsed_sec == 0.0
|
|
|
|
|
|
def test_crawler_today_runs_normally(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""当天日期正常执行 crawler(不受补跑保护影响)。"""
|
|
from datetime import date
|
|
|
|
from scheduler import pipeline
|
|
|
|
called: list[str] = []
|
|
|
|
def _fake_run(cmd, timeout=None): # noqa: ARG001
|
|
called.append(" ".join(cmd))
|
|
return SimpleNamespace(returncode=0)
|
|
|
|
monkeypatch.setattr(pipeline.subprocess, "run", _fake_run)
|
|
sr = pipeline.run_step("crawler", date.today().strftime("%Y%m%d"))
|
|
assert sr.success is True
|
|
assert called and "run_crawler" in called[0]
|
|
|
|
|
|
def test_report_step_uses_date_str(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""report 步骤必须使用传入的 date_str,而非 date.today()(P0-4)。"""
|
|
from scheduler import pipeline, reporter
|
|
|
|
received: list[str] = []
|
|
|
|
def _fake_generate(day_str, upload=True): # noqa: ARG001
|
|
received.append(day_str)
|
|
return 12345
|
|
|
|
monkeypatch.setattr(reporter, "generate_report", _fake_generate)
|
|
sr = pipeline.run_step("report", "20260101")
|
|
assert sr.success is True
|
|
assert received == ["20260101"]
|
|
|
|
|
|
def test_pipeline_backfill_skips_crawler_keeps_rest(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""补跑历史日期:全链路中仅 crawler 跳过,其余步骤照常执行。"""
|
|
from scheduler import pipeline
|
|
|
|
calls: list[str] = []
|
|
|
|
def _fake_run(cmd, timeout=None): # noqa: ARG001
|
|
name = next(c.split(".")[-1] for c in cmd if "scripts.run_" in c)
|
|
calls.append(name)
|
|
return SimpleNamespace(returncode=0)
|
|
|
|
monkeypatch.setattr(pipeline.subprocess, "run", _fake_run)
|
|
state_path = tmp_path / "state.json"
|
|
result = pipeline.run_pipeline(
|
|
"20260101",
|
|
steps=["crawler", "extractor", "dedup"],
|
|
state_path=state_path,
|
|
)
|
|
# crawler 跳过(未执行子进程),extractor/dedup 正常
|
|
assert calls == ["run_extractor", "run_dedup"]
|
|
assert all(s.success for s in result.steps)
|
|
state = pipeline._load_pipeline_state(state_path)
|
|
assert state["20260101"]["crawler"]["status"] == "ok"
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# P1-2: 默认步骤不含 cninfo 独立管道 / P1-3: 调度时区统一
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_default_news_steps_excludes_cninfo_and_report() -> None:
|
|
"""新闻链路默认步骤不含 report 与 cninfo 三步(P1-2)。"""
|
|
from scheduler import DEFAULT_NEWS_STEPS
|
|
|
|
assert "report" not in DEFAULT_NEWS_STEPS
|
|
assert "cninfo_crawl" not in DEFAULT_NEWS_STEPS
|
|
assert "cninfo_extract" not in DEFAULT_NEWS_STEPS
|
|
assert "cninfo_pdf" not in DEFAULT_NEWS_STEPS
|
|
# 新闻链路核心步骤齐全
|
|
assert {"crawler", "xwlb", "extractor", "dedup", "llm",
|
|
"embedding", "qdrant"} <= set(DEFAULT_NEWS_STEPS)
|
|
|
|
|
|
def test_run_pipeline_default_steps_excludes_cninfo(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""run_pipeline 默认 steps 不含 cninfo(只跑新闻链路)。"""
|
|
from scheduler import pipeline
|
|
from scheduler.timeutil import today_str
|
|
|
|
calls: list[str] = []
|
|
|
|
def _fake_run(cmd, timeout=None): # noqa: ARG001
|
|
name = next(c.split(".")[-1] for c in cmd if "scripts.run_" in c)
|
|
calls.append(name)
|
|
return SimpleNamespace(returncode=0)
|
|
|
|
monkeypatch.setattr(pipeline.subprocess, "run", _fake_run)
|
|
state_path = tmp_path / "state.json"
|
|
today = today_str() # 今天:避免 crawler 补跑保护跳过
|
|
result = pipeline.run_pipeline(
|
|
today, state_path=state_path,
|
|
)
|
|
assert result.all_success is True
|
|
assert calls == ["run_crawler", "run_xwlb", "run_extractor", "run_dedup",
|
|
"run_event_extraction", "run_embedding", "run_qdrant_ingest"]
|
|
assert "cninfo" not in " ".join(calls)
|
|
|
|
|
|
def test_scheduler_daemon_backfill_uses_default_news_steps(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""守护进程补跑 steps 使用 DEFAULT_NEWS_STEPS(不含 cninfo 三步)。"""
|
|
import inspect
|
|
|
|
import scripts.run_scheduler as rs
|
|
from scheduler import DEFAULT_NEWS_STEPS
|
|
|
|
# 验证补跑逻辑引用的常量:通过源码断言 + 常量内容双重保证
|
|
src = inspect.getsource(rs._daemon)
|
|
assert "list(DEFAULT_NEWS_STEPS)" in src
|
|
assert "cninfo_crawl" not in DEFAULT_NEWS_STEPS
|
|
|
|
|
|
def test_today_str_matches_schedule_timezone(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""today_str 按调度时区返回 YYYYMMDD(P1-3)。"""
|
|
from datetime import datetime
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from scheduler.timeutil import today_str
|
|
|
|
s = today_str()
|
|
assert len(s) == 8 and s.isdigit()
|
|
# 与调度时区(默认 Asia/Shanghai)当前日期一致
|
|
expect = datetime.now(ZoneInfo("Asia/Shanghai")).strftime("%Y%m%d")
|
|
assert s == expect
|
|
|
|
|
|
def test_today_str_respects_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""SCHEDULE_TZ 环境变量可覆盖调度时区(P1-3)。"""
|
|
from datetime import datetime
|
|
|
|
from scheduler.timeutil import today_str
|
|
|
|
# 覆盖为 UTC 后,today_str 应返回 UTC 日期(而非默认 Asia/Shanghai)
|
|
monkeypatch.setenv("SCHEDULE_TZ", "UTC")
|
|
s = today_str()
|
|
expect = datetime.now(UTC).strftime("%Y%m%d")
|
|
assert s == expect
|
|
|
|
|
|
def test_once_uses_today_str_when_date_missing(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""--once 不带 --date 时按调度时区取今天(P1-3)。"""
|
|
import scripts.run_scheduler as rs
|
|
|
|
captured: dict[str, str] = {}
|
|
|
|
def _fake_pipeline(date_str, steps=None, resume=False): # noqa: ARG001
|
|
captured["date"] = date_str
|
|
|
|
monkeypatch.setattr(rs, "run_pipeline", _fake_pipeline)
|
|
from types import SimpleNamespace
|
|
|
|
args = SimpleNamespace(steps=None, resume=False, date=None)
|
|
rc = rs._once(args)
|
|
assert rc == 0
|
|
assert captured["date"] == rs.today_str()
|