docs: 文档清理与重构 — 统一为 3 个核心文档
- 删除 5 个过时/残留文档(project_plan/agent_prompt/optimization_plan/report_db_design/deploy/README) - 新建 docs/architecture.md(项目架构:11 包职责+数据模型+配置+产物) - 重写 docs/user-guide.md(CLI 全量+增量/断点续跑+MCP+FAQ) - 重写 README.md(精简入口+文档索引) - 更新 continuation.md(追加本次记录) - 更新 .gitignore(排除 data/* 运行产物)
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""tests 包标记。"""
|
||||
@@ -0,0 +1,34 @@
|
||||
"""pytest 共享 fixture。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_sources_yaml(tmp_path: Path) -> Path:
|
||||
"""生成最小可用 sources.yaml 用于测试。"""
|
||||
content = """
|
||||
settings:
|
||||
concurrency: 2
|
||||
retry_max_attempts: 2
|
||||
retry_min_wait_sec: 0.01
|
||||
retry_max_wait_sec: 0.02
|
||||
headless: true
|
||||
output_root: data/raw
|
||||
|
||||
sources:
|
||||
- id: testsrc
|
||||
name: 测试源
|
||||
enabled: true
|
||||
homepage: https://example.com/list
|
||||
article_url_pattern: '^https://example\\.com/article/\\d+$'
|
||||
js_render: false
|
||||
page_timeout_ms: 5000
|
||||
max_articles_per_run: 5
|
||||
"""
|
||||
p = tmp_path / "sources.yaml"
|
||||
p.write_text(content, encoding="utf-8")
|
||||
return p
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,93 @@
|
||||
"""统一 CLI 测试。
|
||||
|
||||
验证子命令路由 + argparse 解析正确。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
from a_share_cli.main import main
|
||||
|
||||
|
||||
def _run(args: str) -> int:
|
||||
with patch.object(sys, "argv", ["a-share", *args.split()]):
|
||||
try:
|
||||
return main()
|
||||
except SystemExit as e:
|
||||
return e.code if isinstance(e.code, int) else 1
|
||||
|
||||
|
||||
def test_no_args_shows_help() -> None:
|
||||
rc = _run("")
|
||||
assert rc == 0
|
||||
|
||||
|
||||
def test_crawl_default() -> None:
|
||||
with patch("a_share_cli.main.cmd_crawl", return_value=0) as mock:
|
||||
_run("crawl")
|
||||
mock.assert_called_once()
|
||||
|
||||
|
||||
def test_crawl_with_source() -> None:
|
||||
with patch("a_share_cli.main.cmd_crawl", return_value=0) as mock:
|
||||
_run("crawl --source cls")
|
||||
args = mock.call_args[0][0]
|
||||
assert args.source == "cls"
|
||||
|
||||
|
||||
def test_extract_with_date() -> None:
|
||||
with patch("a_share_cli.main.cmd_extract", return_value=0) as mock:
|
||||
_run("extract --date 20260616 --source sina")
|
||||
args = mock.call_args[0][0]
|
||||
assert args.date == "20260616"
|
||||
assert args.source == "sina"
|
||||
|
||||
|
||||
def test_events_with_provider() -> None:
|
||||
with patch("a_share_cli.main.cmd_events", return_value=0) as mock:
|
||||
_run("events --provider qwen --limit 5")
|
||||
args = mock.call_args[0][0]
|
||||
assert args.provider == "qwen"
|
||||
assert args.limit == 5
|
||||
|
||||
|
||||
def test_search_default() -> None:
|
||||
with patch("a_share_cli.main.cmd_search", return_value=0) as mock:
|
||||
_run("search 宁德时代")
|
||||
args = mock.call_args[0][0]
|
||||
assert args.query == "宁德时代"
|
||||
assert args.top == 10
|
||||
|
||||
|
||||
def test_search_with_filters() -> None:
|
||||
with patch("a_share_cli.main.cmd_search", return_value=0) as mock:
|
||||
_run("search 芯片 --source cls --sentiment positive --min-importance 3 --top 5")
|
||||
args = mock.call_args[0][0]
|
||||
assert args.query == "芯片"
|
||||
assert args.source == "cls"
|
||||
assert args.sentiment == "positive"
|
||||
assert args.min_importance == 3
|
||||
assert args.top == 5
|
||||
|
||||
|
||||
def test_search_with_stock() -> None:
|
||||
with patch("a_share_cli.main.cmd_search", return_value=0) as mock:
|
||||
_run("search 重大合同 --stock 300750.sz")
|
||||
args = mock.call_args[0][0]
|
||||
assert args.stock == "300750.sz"
|
||||
|
||||
|
||||
def test_pipeline_once() -> None:
|
||||
with patch("a_share_cli.main.cmd_pipeline", return_value=0) as mock:
|
||||
_run("pipeline --once --steps crawler,extractor")
|
||||
args = mock.call_args[0][0]
|
||||
assert args.once is True
|
||||
assert args.steps == "crawler,extractor"
|
||||
|
||||
|
||||
def test_status() -> None:
|
||||
with patch("a_share_cli.main.cmd_status", return_value=0) as mock:
|
||||
_run("status")
|
||||
mock.assert_called_once()
|
||||
@@ -0,0 +1,333 @@
|
||||
"""M1 抓取模块单元测试。
|
||||
|
||||
不依赖真实网络:用 mock 替换 Crawl4AI 的 arun。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import date
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from crawler import engine
|
||||
from crawler.config import load_crawler_config
|
||||
from crawler.engine import (
|
||||
crawl_url_with_retry,
|
||||
extract_article_links,
|
||||
)
|
||||
from crawler.models import CrawlerSettings, CrawlResult, CrawlStage, SourceConfig
|
||||
from crawler.storage import build_output_dir, save_result, url_hash
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 配置加载
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_load_crawler_config_ok(sample_sources_yaml: Path) -> None:
|
||||
cfg = load_crawler_config(sample_sources_yaml)
|
||||
assert cfg.settings.concurrency == 2
|
||||
assert len(cfg.sources) == 1
|
||||
assert cfg.sources[0].id == "testsrc"
|
||||
assert cfg.enabled_sources()[0].id == "testsrc"
|
||||
|
||||
|
||||
def test_load_crawler_config_real_sources_yaml() -> None:
|
||||
"""项目内置的 configs/sources.yaml 必须可解析(回归保护)。"""
|
||||
real = Path("configs/sources.yaml")
|
||||
if not real.is_file():
|
||||
pytest.skip("configs/sources.yaml 未生成,跳过")
|
||||
cfg = load_crawler_config(real)
|
||||
# 验收标准: 至少 5 个启用源
|
||||
assert len(cfg.enabled_sources()) >= 5, "启用源应至少 5 个(M1 验收标准)"
|
||||
assert cfg.settings.concurrency == 3, "用户决策: 并发上限 3"
|
||||
|
||||
|
||||
def test_load_crawler_config_missing_file(tmp_path: Path) -> None:
|
||||
with pytest.raises(FileNotFoundError):
|
||||
load_crawler_config(tmp_path / "nonexistent.yaml")
|
||||
|
||||
|
||||
def test_source_id_validation() -> None:
|
||||
with pytest.raises(ValueError):
|
||||
SourceConfig(
|
||||
id="Bad-ID", # 含大写与连字符
|
||||
name="x",
|
||||
homepage="https://example.com",
|
||||
article_url_pattern="^.*$",
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 链接抽取
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _src(**kwargs: Any) -> SourceConfig:
|
||||
base: dict[str, Any] = {
|
||||
"id": "testsrc",
|
||||
"name": "测试",
|
||||
"homepage": "https://example.com/list",
|
||||
"article_url_pattern": r"^https://example\.com/article/\d+$",
|
||||
"js_render": False,
|
||||
"max_articles_per_run": 10,
|
||||
}
|
||||
base.update(kwargs)
|
||||
return SourceConfig(**base)
|
||||
|
||||
|
||||
def test_extract_article_links_basic() -> None:
|
||||
html = """
|
||||
<html><body>
|
||||
<a href="/article/123">文章一</a>
|
||||
<a href="https://example.com/article/456">文章二</a>
|
||||
<a href="https://other.com/article/789">外站</a>
|
||||
<a href="/about">关于</a>
|
||||
<a href="javascript:void(0)">JS</a>
|
||||
<a href="/article/123">文章一(重复)</a>
|
||||
</body></html>
|
||||
"""
|
||||
links = extract_article_links(html, "https://example.com/list", _src())
|
||||
urls = [link.url for link in links]
|
||||
assert urls == [
|
||||
"https://example.com/article/123",
|
||||
"https://example.com/article/456",
|
||||
]
|
||||
assert links[0].anchor_text == "文章一"
|
||||
|
||||
|
||||
def test_extract_article_links_respects_max() -> None:
|
||||
html_parts = [
|
||||
f'<a href="https://example.com/article/{i}">a{i}</a>' for i in range(20)
|
||||
]
|
||||
html = "<html><body>" + "".join(html_parts) + "</body></html>"
|
||||
src = _src(max_articles_per_run=3)
|
||||
links = extract_article_links(html, "https://example.com/list", src)
|
||||
assert len(links) == 3
|
||||
|
||||
|
||||
def test_extract_article_links_strips_fragment() -> None:
|
||||
html = '<a href="https://example.com/article/1#section">x</a>'
|
||||
links = extract_article_links(html, "https://example.com/list", _src())
|
||||
assert links[0].url == "https://example.com/article/1"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 重试机制
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class _FakeC4Result:
|
||||
"""模拟 Crawl4AI 的返回对象。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
success: bool = True,
|
||||
html: str = "<html>ok</html>",
|
||||
markdown: str = "ok",
|
||||
status_code: int = 200,
|
||||
error_message: str | None = None,
|
||||
) -> None:
|
||||
self.success = success
|
||||
self.html = html
|
||||
self.markdown = markdown
|
||||
self.status_code = status_code
|
||||
self.error_message = error_message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_succeeds_on_third_attempt() -> None:
|
||||
"""前两次失败,第三次成功;返回的 attempts 应为 3。"""
|
||||
fake_crawler = AsyncMock()
|
||||
fake_crawler.arun = AsyncMock(
|
||||
side_effect=[
|
||||
_FakeC4Result(success=False, html="", error_message="boom-1"),
|
||||
_FakeC4Result(success=False, html="", error_message="boom-2"),
|
||||
_FakeC4Result(success=True),
|
||||
]
|
||||
)
|
||||
|
||||
settings = CrawlerSettings(
|
||||
concurrency=1,
|
||||
retry_max_attempts=3,
|
||||
retry_min_wait_sec=0.0,
|
||||
retry_max_wait_sec=0.0,
|
||||
)
|
||||
sem = asyncio.Semaphore(1)
|
||||
src = _src()
|
||||
res = await crawl_url_with_retry(
|
||||
crawler=fake_crawler,
|
||||
url="https://example.com/article/1",
|
||||
source=src,
|
||||
stage=CrawlStage.ARTICLE,
|
||||
settings=settings,
|
||||
semaphore=sem,
|
||||
)
|
||||
assert res.success is True
|
||||
assert res.attempts == 3
|
||||
assert fake_crawler.arun.await_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_gives_up_after_max() -> None:
|
||||
fake_crawler = AsyncMock()
|
||||
fake_crawler.arun = AsyncMock(
|
||||
return_value=_FakeC4Result(success=False, html="", error_message="nope")
|
||||
)
|
||||
settings = CrawlerSettings(
|
||||
concurrency=1,
|
||||
retry_max_attempts=2,
|
||||
retry_min_wait_sec=0.0,
|
||||
retry_max_wait_sec=0.0,
|
||||
)
|
||||
sem = asyncio.Semaphore(1)
|
||||
res = await crawl_url_with_retry(
|
||||
crawler=fake_crawler,
|
||||
url="https://example.com/article/1",
|
||||
source=_src(),
|
||||
stage=CrawlStage.ARTICLE,
|
||||
settings=settings,
|
||||
semaphore=sem,
|
||||
)
|
||||
assert res.success is False
|
||||
assert res.attempts == 2
|
||||
assert fake_crawler.arun.await_count == 2
|
||||
assert res.error == "nope"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exception_is_swallowed_and_retried() -> None:
|
||||
"""arun 抛异常应被捕获并触发重试。"""
|
||||
fake_crawler = AsyncMock()
|
||||
fake_crawler.arun = AsyncMock(
|
||||
side_effect=[RuntimeError("net down"), _FakeC4Result(success=True)]
|
||||
)
|
||||
settings = CrawlerSettings(
|
||||
concurrency=1, retry_max_attempts=2, retry_min_wait_sec=0.0, retry_max_wait_sec=0.0
|
||||
)
|
||||
sem = asyncio.Semaphore(1)
|
||||
res = await crawl_url_with_retry(
|
||||
crawler=fake_crawler,
|
||||
url="https://example.com/article/1",
|
||||
source=_src(),
|
||||
stage=CrawlStage.ARTICLE,
|
||||
settings=settings,
|
||||
semaphore=sem,
|
||||
)
|
||||
assert res.success is True
|
||||
assert res.attempts == 2
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 存储
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_url_hash_stable() -> None:
|
||||
h1 = url_hash("https://example.com/a")
|
||||
h2 = url_hash("https://example.com/a")
|
||||
h3 = url_hash("https://example.com/b")
|
||||
assert h1 == h2
|
||||
assert h1 != h3
|
||||
assert len(h1) == 16
|
||||
|
||||
|
||||
def test_build_output_dir(tmp_path: Path) -> None:
|
||||
d = build_output_dir(tmp_path, "cls", date(2026, 6, 16))
|
||||
assert d == tmp_path / "cls" / "20260616"
|
||||
|
||||
|
||||
def test_save_result_writes_html_md_and_index(tmp_path: Path) -> None:
|
||||
result = CrawlResult(
|
||||
source_id="cls",
|
||||
stage=CrawlStage.ARTICLE,
|
||||
url="https://example.com/article/1",
|
||||
success=True,
|
||||
status_code=200,
|
||||
title="标题",
|
||||
html="<html>hi</html>",
|
||||
markdown="# hi",
|
||||
)
|
||||
html_path = save_result(result, tmp_path, day=date(2026, 6, 16))
|
||||
assert html_path is not None and html_path.is_file()
|
||||
assert html_path.read_text(encoding="utf-8") == "<html>hi</html>"
|
||||
|
||||
md_path = html_path.with_suffix(".md")
|
||||
assert md_path.is_file()
|
||||
assert md_path.read_text(encoding="utf-8") == "# hi"
|
||||
|
||||
index = html_path.parent / "index.jsonl"
|
||||
assert index.is_file()
|
||||
line = index.read_text(encoding="utf-8").strip()
|
||||
obj = json.loads(line)
|
||||
assert obj["url"] == "https://example.com/article/1"
|
||||
assert obj["success"] is True
|
||||
assert obj["html_file"] == html_path.name
|
||||
assert "html" not in obj # 大字段不应进入元数据
|
||||
|
||||
|
||||
def test_save_result_failure_only_appends_index(tmp_path: Path) -> None:
|
||||
result = CrawlResult(
|
||||
source_id="cls",
|
||||
stage=CrawlStage.ARTICLE,
|
||||
url="https://example.com/article/2",
|
||||
success=False,
|
||||
error="timeout",
|
||||
)
|
||||
html_path = save_result(result, tmp_path, day=date(2026, 6, 16))
|
||||
assert html_path is None
|
||||
index = tmp_path / "cls" / "20260616" / "index.jsonl"
|
||||
assert index.is_file()
|
||||
obj = json.loads(index.read_text(encoding="utf-8").strip())
|
||||
assert obj["success"] is False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Markdown 兼容
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_markdown_text_handles_str() -> None:
|
||||
assert engine._markdown_text("plain") == "plain"
|
||||
|
||||
|
||||
def test_markdown_text_handles_object_raw_markdown() -> None:
|
||||
class _Obj:
|
||||
raw_markdown = "from raw"
|
||||
|
||||
assert engine._markdown_text(_Obj()) == "from raw"
|
||||
|
||||
|
||||
def test_markdown_text_handles_none() -> None:
|
||||
assert engine._markdown_text(None) == ""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 集成测试(默认跳过,需要真实浏览器与网络)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_homepage_crawl_smoke() -> None:
|
||||
"""真实抓取 example.com 烟测,验证端到端可运行。
|
||||
|
||||
运行: uv run pytest -m integration
|
||||
"""
|
||||
from crawler import crawl_all
|
||||
from crawler.models import CrawlerConfig
|
||||
|
||||
cfg = CrawlerConfig(
|
||||
settings=CrawlerSettings(concurrency=1, retry_max_attempts=1, output_root="data/raw_test"),
|
||||
sources=[
|
||||
SourceConfig(
|
||||
id="example",
|
||||
name="example",
|
||||
homepage="https://example.com/",
|
||||
article_url_pattern=r"^https://www\.iana\.org/.*$",
|
||||
js_render=False,
|
||||
page_timeout_ms=15000,
|
||||
max_articles_per_run=1,
|
||||
)
|
||||
],
|
||||
)
|
||||
results = await crawl_all(cfg, save=False)
|
||||
assert any(r.success for r in results), "example.com 烟测应至少一个成功"
|
||||
@@ -0,0 +1,502 @@
|
||||
"""M3 三层去重模块单元测试。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from dedup import (
|
||||
DEFAULT_HAMMING_THRESHOLD,
|
||||
Deduper,
|
||||
DedupLayer,
|
||||
Fingerprint,
|
||||
FingerprintStore,
|
||||
article_to_fingerprint,
|
||||
content_hash,
|
||||
hamming,
|
||||
normalize_content,
|
||||
simhash64,
|
||||
)
|
||||
from extractor import Article
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# fixtures
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _article(
|
||||
*,
|
||||
url: str = "https://www.cls.cn/detail/1",
|
||||
url_hash: str = "abc1234567890000",
|
||||
source_id: str = "cls",
|
||||
title: str = "宁德时代发布新一代麒麟电池",
|
||||
content: str = (
|
||||
"宁德时代今日正式发布了新一代麒麟电池产品，能量密度达到 255 Wh/kg，"
|
||||
"显著优于上一代产品。该电池将于2026年第三季度量产。"
|
||||
),
|
||||
publish_time: datetime | None = datetime(2026, 6, 16, 10, 0),
|
||||
) -> Article:
|
||||
return Article(
|
||||
source_id=source_id,
|
||||
url=url,
|
||||
url_hash=url_hash,
|
||||
title=title,
|
||||
content=content,
|
||||
publish_time=publish_time,
|
||||
word_count=len(content),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tmp_db(tmp_path: Path) -> Path:
|
||||
return tmp_path / "fp.sqlite3"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# hasher
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_normalize_content_strips_punct_and_whitespace() -> None:
|
||||
norm = normalize_content("你好, 世界!\n这是 中文。")
|
||||
assert norm == "你好世界这是中文"
|
||||
|
||||
|
||||
def test_normalize_content_handles_empty() -> None:
|
||||
assert normalize_content("") == ""
|
||||
assert normalize_content(" \n\t ") == ""
|
||||
|
||||
|
||||
def test_content_hash_deterministic_and_punct_invariant() -> None:
|
||||
a = "今天天气很好。"
|
||||
b = "今天，天气，很好！！！"
|
||||
assert content_hash(a) == content_hash(b)
|
||||
|
||||
|
||||
def test_content_hash_differs_for_different_text() -> None:
|
||||
assert content_hash("今天天气很好") != content_hash("今天天气不好")
|
||||
|
||||
|
||||
def test_simhash_identical_text_same_value() -> None:
|
||||
text = "宁德时代发布新一代麒麟电池产品 能量密度大幅提升"
|
||||
assert simhash64(text) == simhash64(text)
|
||||
|
||||
|
||||
def test_simhash_minor_changes_close_distance() -> None:
|
||||
"""长文本(贴近真实新闻)的轻度改写,hamming 距离应在阈值内。"""
|
||||
base = (
|
||||
"宁德时代今日正式发布新一代麒麟电池产品，能量密度达到 255 瓦时每公斤，"
|
||||
"显著优于上一代产品。该电池将于 2026 年第三季度量产，首批应用于多款新能源汽车。"
|
||||
"公司股价应声上涨 5.2%，分析师认为这将进一步巩固宁德时代在全球动力电池领域的领先地位。"
|
||||
) * 2
|
||||
rewritten = "财联社讯：" + base + "(完)"
|
||||
d = hamming(simhash64(base), simhash64(rewritten))
|
||||
assert d <= DEFAULT_HAMMING_THRESHOLD, f"长文本前后加标识汉明距离 {d} 不应超过阈值"
|
||||
|
||||
|
||||
def test_simhash_unrelated_text_far_distance() -> None:
|
||||
"""完全不相关的两段长文本汉明距离应远大于阈值。"""
|
||||
a = "宁德时代发布新一代麒麟电池产品，能量密度达到255瓦时每公斤。" * 3
|
||||
b = "美联储宣布维持利率不变，市场普遍预期下次会议将开启降息周期。" * 3
|
||||
d = hamming(simhash64(a), simhash64(b))
|
||||
assert d > DEFAULT_HAMMING_THRESHOLD * 2
|
||||
|
||||
|
||||
def test_simhash_empty_returns_zero() -> None:
|
||||
assert simhash64("") == 0
|
||||
assert simhash64(" ") == 0
|
||||
|
||||
|
||||
def test_hamming_basics() -> None:
|
||||
assert hamming(0, 0) == 0
|
||||
assert hamming(0xFF, 0x00) == 8
|
||||
assert hamming(0xFF00FF00, 0x00FF00FF) == 32
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# FingerprintStore
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_store_upsert_and_get(tmp_db: Path) -> None:
|
||||
fp = Fingerprint(
|
||||
url_hash="hash1",
|
||||
content_hash="ch1",
|
||||
simhash=0xDEADBEEFCAFEBABE,
|
||||
source_id="cls",
|
||||
url="https://x/1",
|
||||
title="A",
|
||||
publish_date="2026-06-16",
|
||||
)
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
store.upsert(fp)
|
||||
got = store.get_by_url_hash("hash1")
|
||||
assert got is not None
|
||||
assert got.content_hash == "ch1"
|
||||
assert got.simhash == 0xDEADBEEFCAFEBABE
|
||||
assert got.publish_date == "2026-06-16"
|
||||
|
||||
|
||||
def test_store_upsert_replaces_existing(tmp_db: Path) -> None:
|
||||
base = Fingerprint(
|
||||
url_hash="h",
|
||||
content_hash="ch1",
|
||||
simhash=1,
|
||||
source_id="cls",
|
||||
url="u",
|
||||
title="t",
|
||||
)
|
||||
updated = base.model_copy(update={"content_hash": "ch2", "simhash": 999})
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
store.upsert(base)
|
||||
store.upsert(updated)
|
||||
got = store.get_by_url_hash("h")
|
||||
assert got is not None
|
||||
assert got.content_hash == "ch2"
|
||||
assert got.simhash == 999
|
||||
assert store.count() == 1
|
||||
|
||||
|
||||
def test_store_find_by_content_hash(tmp_db: Path) -> None:
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
store.upsert(Fingerprint(
|
||||
url_hash="h1", content_hash="ch", simhash=0,
|
||||
source_id="cls", url="u1", title="t1"
|
||||
))
|
||||
assert store.find_by_content_hash("ch") is not None
|
||||
assert store.find_by_content_hash("nope") is None
|
||||
|
||||
|
||||
def test_store_candidates_within_window(tmp_db: Path) -> None:
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
for d, h in [("2026-05-01", "old"), ("2026-06-15", "near"), ("2026-07-30", "far")]:
|
||||
store.upsert(Fingerprint(
|
||||
url_hash=h, content_hash=h, simhash=0,
|
||||
source_id="cls", url=f"u/{h}", title=h, publish_date=d,
|
||||
))
|
||||
cands = store.candidates_for_simhash("2026-06-16", window_days=7)
|
||||
url_hashes = sorted(c.url_hash for c in cands)
|
||||
assert url_hashes == ["near"]
|
||||
|
||||
|
||||
def test_store_candidates_no_date_returns_all(tmp_db: Path) -> None:
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
store.upsert(Fingerprint(
|
||||
url_hash="h1", content_hash="c1", simhash=0,
|
||||
source_id="cls", url="u", title="t", publish_date=None
|
||||
))
|
||||
cands = store.candidates_for_simhash(None, 30)
|
||||
assert len(cands) == 1
|
||||
|
||||
|
||||
def test_store_simhash_handles_high_bit(tmp_db: Path) -> None:
|
||||
"""64 位 SimHash 高位为 1 时,hex 存取应保持无符号。"""
|
||||
high = (1 << 63) | 0x1234
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
store.upsert(Fingerprint(
|
||||
url_hash="h", content_hash="c", simhash=high,
|
||||
source_id="cls", url="u", title="t",
|
||||
))
|
||||
got = store.get_by_url_hash("h")
|
||||
assert got is not None
|
||||
assert got.simhash == high
|
||||
|
||||
|
||||
def test_store_count_by_source(tmp_db: Path) -> None:
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
for i, src in enumerate(["cls", "cls", "sina"]):
|
||||
store.upsert(Fingerprint(
|
||||
url_hash=f"h{i}", content_hash=f"c{i}", simhash=i,
|
||||
source_id=src, url=f"u{i}", title=f"t{i}",
|
||||
))
|
||||
counts = store.count_by_source()
|
||||
assert counts == {"cls": 2, "sina": 1}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Deduper - 三层判重
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_dedup_first_article_is_unique(tmp_db: Path) -> None:
|
||||
art = _article()
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
result = d.ingest(art)
|
||||
assert not result.is_duplicate
|
||||
assert result.matched_layer is None
|
||||
assert d.stats().total == 1
|
||||
|
||||
|
||||
def test_dedup_layer1_url_hash(tmp_db: Path) -> None:
|
||||
"""同一 url_hash 直接命中 L1。"""
|
||||
a1 = _article()
|
||||
a2 = _article() # 同 url_hash 同 url
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
d.ingest(a1)
|
||||
result = d.ingest(a2)
|
||||
assert result.is_duplicate
|
||||
assert result.matched_layer == DedupLayer.URL
|
||||
assert d.stats().total == 1, "L1 命中应不写入新指纹"
|
||||
|
||||
|
||||
def test_dedup_layer2_content_hash(tmp_db: Path) -> None:
|
||||
"""url 不同但 content 完全一致 -> L2。"""
|
||||
a1 = _article(url="https://a.com/1", url_hash="hash1aaaaaaaaaaa")
|
||||
a2 = _article(url="https://b.com/2", url_hash="hash2bbbbbbbbbbb")
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
d.ingest(a1)
|
||||
result = d.ingest(a2)
|
||||
assert result.is_duplicate
|
||||
assert result.matched_layer == DedupLayer.CONTENT
|
||||
assert result.matched_url_hash == "hash1aaaaaaaaaaa"
|
||||
|
||||
|
||||
def test_dedup_layer2_punctuation_difference_still_caught(tmp_db: Path) -> None:
|
||||
"""标点/空白差异不应阻止 L2 命中(normalize_content 应剥离)。"""
|
||||
base = "今天天气很好我们去公园散步"
|
||||
a1 = _article(
|
||||
url="https://a/1", url_hash="aaaa", content="今天天气很好。我们去公园散步！"
|
||||
)
|
||||
a2 = _article(
|
||||
url="https://b/2", url_hash="bbbb", content="今天天气，很好；我们去公园 散步！！"
|
||||
)
|
||||
assert content_hash(a1.content) == content_hash(a2.content)
|
||||
assert normalize_content(a1.content) == base
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
d.ingest(a1)
|
||||
result = d.ingest(a2)
|
||||
assert result.matched_layer == DedupLayer.CONTENT
|
||||
|
||||
|
||||
def test_dedup_layer3_simhash_minor_rewrite(tmp_db: Path) -> None:
|
||||
"""长文本 + 转载前后缀,落入 SimHash 层(贴近真实跨源转载场景)。"""
|
||||
long_body = (
|
||||
"宁德时代今日正式发布新一代麒麟电池产品，能量密度达到 255 瓦时每公斤，"
|
||||
"显著优于上一代产品。该电池将于 2026 年第三季度量产，首批应用于多款新能源汽车。"
|
||||
"公司股价应声上涨 5.2%，分析师认为这将进一步巩固宁德时代在全球动力电池领域的领先地位。"
|
||||
) * 2
|
||||
rewritten = "财联社讯：" + long_body + "(完)"
|
||||
a1 = _article(url="https://a/1", url_hash="aaaaa", content=long_body)
|
||||
a2 = _article(url="https://b/2", url_hash="bbbbb", content=rewritten)
|
||||
# 必要前提:content_hash 不同(否则会被 L2 截胡)
|
||||
assert content_hash(a1.content) != content_hash(a2.content)
|
||||
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
d.ingest(a1)
|
||||
result = d.ingest(a2)
|
||||
assert result.is_duplicate
|
||||
assert result.matched_layer == DedupLayer.SIMHASH
|
||||
assert result.hamming_distance is not None
|
||||
assert result.hamming_distance <= DEFAULT_HAMMING_THRESHOLD
|
||||
|
||||
|
||||
def test_dedup_layer3_unrelated_articles_kept(tmp_db: Path) -> None:
|
||||
a1 = _article(url="https://a/1", url_hash="aaaaa",
|
||||
content="宁德时代发布新一代麒麟电池产品，能量密度达到 255 瓦时每公斤。" * 5)
|
||||
a2 = _article(url="https://b/2", url_hash="bbbbb",
|
||||
content="美联储宣布维持联邦基金利率不变，市场预期下次会议将开启降息。" * 5,
|
||||
title="美联储利率决议")
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
d.ingest(a1)
|
||||
result = d.ingest(a2)
|
||||
assert not result.is_duplicate
|
||||
assert d.stats().total == 2
|
||||
|
||||
|
||||
def test_dedup_layer3_outside_time_window_kept(tmp_db: Path) -> None:
|
||||
"""SimHash 相近,但 publish_date 距离过远(> 30 天)不去重。"""
|
||||
body = (
|
||||
"宁德时代今日正式发布新一代麒麟电池产品，能量密度达到 255 瓦时每公斤，"
|
||||
"显著优于上一代产品。该电池将于第三季度量产，首批应用于多款新能源汽车。" * 2
|
||||
)
|
||||
a1 = _article(
|
||||
url="https://a/1", url_hash="aaaa1", content=body,
|
||||
publish_time=datetime(2026, 1, 1, 9, 0),
|
||||
)
|
||||
a2 = _article(
|
||||
url="https://b/2", url_hash="bbbb2", content=body[3:], # 微改 -> 走 L3
|
||||
publish_time=datetime(2026, 6, 16, 9, 0),
|
||||
)
|
||||
assert content_hash(a1.content) != content_hash(a2.content)
|
||||
with Deduper(db_path=tmp_db, time_window_days=30) as d:
|
||||
d.ingest(a1)
|
||||
result = d.ingest(a2)
|
||||
assert not result.is_duplicate, "时间窗口外不应命中 SimHash"
|
||||
|
||||
|
||||
def test_dedup_threshold_zero_only_exact_simhash(tmp_db: Path) -> None:
|
||||
"""阈值 0 -> 仅当 SimHash 完全相同才视为重复(且会先被 L2 拦截)。"""
|
||||
a1 = _article(url="https://a/1", url_hash="aaaa1",
|
||||
content="宁德时代发布新一代麒麟电池产品 能量密度大幅提升")
|
||||
a2 = _article(url="https://b/2", url_hash="bbbb2",
|
||||
content="财联社讯 宁德时代今天发布了新一代麒麟电池 能量密度提升明显")
|
||||
with Deduper(db_path=tmp_db, simhash_threshold=0) as d:
|
||||
d.ingest(a1)
|
||||
result = d.ingest(a2)
|
||||
# 两段相似但不同的文本,阈值 0 时不应判重
|
||||
assert not result.is_duplicate
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Deduper - check 不写入
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_check_does_not_write(tmp_db: Path) -> None:
|
||||
art = _article()
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
result = d.check(art)
|
||||
assert not result.is_duplicate
|
||||
assert d.stats().total == 0 # check 不应入库
|
||||
|
||||
|
||||
def test_article_to_fingerprint_fields() -> None:
|
||||
art = _article()
|
||||
fp = article_to_fingerprint(art)
|
||||
assert fp.url_hash == art.url_hash
|
||||
assert fp.simhash == simhash64(art.content)
|
||||
assert fp.content_hash == content_hash(art.content)
|
||||
assert fp.publish_date == "2026-06-16"
|
||||
|
||||
|
||||
def test_article_to_fingerprint_handles_none_publish_time() -> None:
|
||||
art = _article(publish_time=None)
|
||||
fp = article_to_fingerprint(art)
|
||||
assert fp.publish_date is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# stats
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_stats_aggregates_by_source(tmp_db: Path) -> None:
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
d.ingest(_article(source_id="cls", url="https://cls/1",
|
||||
url_hash="cls0000000000001"))
|
||||
d.ingest(_article(source_id="cls", url="https://cls/2",
|
||||
url_hash="cls0000000000002",
|
||||
content="完全不同的另一篇文章" * 30))
|
||||
d.ingest(_article(source_id="sina", url="https://sina/1",
|
||||
url_hash="sina000000000001",
|
||||
content="第三篇 完全不同 主题 美联储 利率" * 20))
|
||||
stats = d.stats()
|
||||
assert stats.total == 3
|
||||
assert stats.by_source == {"cls": 2, "sina": 1}
|
||||
assert stats.earliest is not None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 多源记录(source_ids)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_fingerprint_source_ids_default_to_source() -> None:
|
||||
"""source_ids 未显式给定时,自动包含主源 source_id。"""
|
||||
fp = Fingerprint(
|
||||
url_hash="h", content_hash="c", simhash=0,
|
||||
source_id="cls", url="u", title="t",
|
||||
)
|
||||
assert fp.source_ids == ["cls"]
|
||||
|
||||
|
||||
def test_fingerprint_source_ids_keeps_main_source_first() -> None:
|
||||
"""source_ids 无论怎么传,主源 source_id 始终居首且去重。"""
|
||||
fp = Fingerprint(
|
||||
url_hash="h", content_hash="c", simhash=0,
|
||||
source_id="cls", url="u", title="t",
|
||||
source_ids=["sina", "cls", "eastmoney", "sina"],
|
||||
)
|
||||
assert fp.source_ids[0] == "cls"
|
||||
assert len(fp.source_ids) == len(set(fp.source_ids)) # 无重复
|
||||
|
||||
|
||||
def test_store_persists_source_ids(tmp_db: Path) -> None:
|
||||
fp = Fingerprint(
|
||||
url_hash="h", content_hash="c", simhash=0,
|
||||
source_id="cls", url="u", title="t",
|
||||
source_ids=["cls", "sina", "eastmoney"],
|
||||
)
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
store.upsert(fp)
|
||||
got = store.get_by_url_hash("h")
|
||||
assert got is not None
|
||||
assert got.source_ids == ["cls", "sina", "eastmoney"]
|
||||
|
||||
|
||||
def test_store_migrates_old_schema_without_source_ids(tmp_db: Path) -> None:
|
||||
"""旧库(无 source_ids 列)打开时应自动迁移,旧数据回退为 [source_id]。"""
|
||||
import sqlite3
|
||||
|
||||
conn = sqlite3.connect(tmp_db)
|
||||
conn.executescript(
|
||||
"CREATE TABLE fingerprints ("
|
||||
" url_hash TEXT PRIMARY KEY, content_hash TEXT NOT NULL, simhash_hex TEXT NOT NULL,"
|
||||
" source_id TEXT NOT NULL, url TEXT NOT NULL, title TEXT NOT NULL,"
|
||||
" publish_date TEXT, ingested_at TEXT NOT NULL);"
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO fingerprints VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
("old1", "ch1", "0000000000000000", "cls", "u1", "t1", "2026-06-01", "2026-06-01T00:00:00"),
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
with FingerprintStore(tmp_db) as store:
|
||||
got = store.get_by_url_hash("old1")
|
||||
assert got is not None
|
||||
assert got.source_ids == ["cls"] # 迁移后回退主源
|
||||
# 迁移后可正常写入多源
|
||||
store.upsert(Fingerprint(
|
||||
url_hash="new1", content_hash="c2", simhash=1,
|
||||
source_id="sina", url="u2", title="t2",
|
||||
source_ids=["sina", "cls"],
|
||||
))
|
||||
assert store.get_by_url_hash("new1") is not None # type: ignore[union-attr]
|
||||
|
||||
|
||||
def test_ingest_merges_sources_on_duplicate(tmp_db: Path) -> None:
|
||||
"""同一内容被多个源发布时,重复文章的来源并入唯一新闻指纹。"""
|
||||
body = "宁德时代今日发布新一代麒麟电池,能量密度 255Wh/kg。" * 4
|
||||
a1 = _article(source_id="cls", url="https://cls/a", url_hash="aaaa111111111111",
|
||||
content=body)
|
||||
a2 = _article(source_id="sina", url="https://sina/b", url_hash="bbbb222222222222",
|
||||
content=body)
|
||||
a3 = _article(source_id="eastmoney", url="https://em/c", url_hash="cccc333333333333",
|
||||
content=body)
|
||||
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
r1 = d.ingest(a1)
|
||||
assert not r1.is_duplicate
|
||||
r2 = d.ingest(a2)
|
||||
assert r2.is_duplicate
|
||||
assert r2.matched_layer == DedupLayer.CONTENT
|
||||
assert r2.matched_source_id == "cls"
|
||||
# 命中后 all_source_ids 立即包含两个源
|
||||
assert r2.all_source_ids == ["cls", "sina"]
|
||||
|
||||
r3 = d.ingest(a3)
|
||||
assert r3.is_duplicate
|
||||
assert r3.all_source_ids == ["cls", "sina", "eastmoney"]
|
||||
|
||||
# 指纹库持久化多源
|
||||
matched = d.store.get_by_url_hash("aaaa111111111111")
|
||||
assert matched is not None
|
||||
assert matched.source_ids == ["cls", "sina", "eastmoney"]
|
||||
assert d.stats().total == 1 # 内容组只算 1 条唯一
|
||||
|
||||
|
||||
def test_check_reports_all_sources_without_writing(tmp_db: Path) -> None:
|
||||
"""check(只读)命中重复时也能看到全部来源,且不写库。"""
|
||||
body = "宁德时代发布新一代麒麟电池产品。" * 6
|
||||
a1 = _article(source_id="cls", url="https://cls/a", url_hash="aaaa111111111111",
|
||||
content=body)
|
||||
a2 = _article(source_id="sina", url="https://sina/b", url_hash="bbbb222222222222",
|
||||
content=body)
|
||||
|
||||
with Deduper(db_path=tmp_db) as d:
|
||||
d.ingest(a1)
|
||||
d.ingest(a2)
|
||||
# 第三次来一篇同样内容的文章,仅 check
|
||||
a3 = _article(source_id="eastmoney", url="https://em/c", url_hash="cccc333333333333",
|
||||
content=body)
|
||||
result = d.check(a3)
|
||||
assert result.is_duplicate
|
||||
assert result.matched_source_id == "cls"
|
||||
assert result.all_source_ids == ["cls", "sina"]
|
||||
assert d.stats().total == 1 # check 不写库
|
||||
@@ -0,0 +1,397 @@
|
||||
"""M5 嵌入模块单元测试。
|
||||
|
||||
不依赖真实 LLM/HuggingFace,所有 provider 调用通过 mock 注入。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from embedding import (
|
||||
DASHSCOPE_BATCH_LIMIT,
|
||||
AsyncEmbeddingProvider,
|
||||
EmbeddingError,
|
||||
EmbeddingProvider,
|
||||
EmbeddingProviderType,
|
||||
EmbeddingResult,
|
||||
compose_text,
|
||||
make_async_provider,
|
||||
make_sync_provider,
|
||||
resolve_provider_type,
|
||||
)
|
||||
from embedding.base import _from_event_dict
|
||||
from embedding.remote import (
|
||||
DashScopeAsyncEmbeddingProvider,
|
||||
DashScopeEmbeddingProvider,
|
||||
_chunked,
|
||||
)
|
||||
from extractor import Article
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# fixtures
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _article(
|
||||
*,
|
||||
url: str = "https://www.cls.cn/detail/1",
|
||||
url_hash: str = "abc1234567890000",
|
||||
title: str = "宁德时代签订 100GWh 长期供货协议",
|
||||
content: str = "宁德时代与某车企签 5 年 100GWh 协议,涉及金额超 1500 亿。" * 3,
|
||||
publish_time: datetime | None = datetime(2026, 6, 16, 10, 0),
|
||||
) -> Article:
|
||||
return Article(
|
||||
source_id="cls",
|
||||
url=url,
|
||||
url_hash=url_hash,
|
||||
title=title,
|
||||
content=content,
|
||||
publish_time=publish_time,
|
||||
word_count=len(content),
|
||||
)
|
||||
|
||||
|
||||
def _embedding_response(vectors: list[list[float]]) -> MagicMock:
|
||||
"""构造与 OpenAI SDK 一致的 embeddings.create 返回。"""
|
||||
resp = MagicMock()
|
||||
resp.data = [MagicMock(embedding=v) for v in vectors]
|
||||
return resp
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# compose_text
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_compose_text_basic() -> None:
|
||||
art = _article()
|
||||
text = compose_text(art)
|
||||
assert text.startswith("标题:")
|
||||
assert "正文:" in text
|
||||
assert art.title in text
|
||||
assert art.content[:30] in text
|
||||
|
||||
|
||||
def test_compose_text_with_head_and_summary() -> None:
|
||||
art = _article()
|
||||
text = compose_text(art, head="[sentiment=positive]", summary="一句话摘要")
|
||||
assert "[sentiment=positive]" in text
|
||||
assert "摘要:一句话摘要" in text
|
||||
|
||||
|
||||
def test_compose_text_truncates_overlong() -> None:
|
||||
art = _article(content="字" * 10000)
|
||||
text = compose_text(art, max_chars=500)
|
||||
assert len(text) <= 500
|
||||
|
||||
|
||||
def test_from_event_dict_extracts_head_and_article() -> None:
|
||||
event_obj = {
|
||||
"source_id": "cls",
|
||||
"url": "https://x/1",
|
||||
"url_hash": "h1",
|
||||
"title": "宁德合作",
|
||||
"publish_time": "2026-06-16T10:00:00",
|
||||
"event": {
|
||||
"stock_codes": ["300750.SZ"],
|
||||
"company_names": ["宁德时代"],
|
||||
"industries": ["动力电池"],
|
||||
"sentiment": "positive",
|
||||
"importance": 5,
|
||||
"event_type": "重大合同",
|
||||
"summary": "签订 100GWh 协议",
|
||||
},
|
||||
}
|
||||
article, head, summary = _from_event_dict(event_obj)
|
||||
assert article.url_hash == "h1"
|
||||
assert article.publish_time == datetime(2026, 6, 16, 10, 0, 0)
|
||||
assert "sentiment=positive" in head
|
||||
assert "importance=5" in head
|
||||
assert "300750.SZ" in head
|
||||
assert "宁德时代" in head
|
||||
assert "动力电池" in head
|
||||
assert summary == "签订 100GWh 协议"
|
||||
|
||||
|
||||
def test_from_event_dict_handles_missing_publish_time() -> None:
|
||||
article, _, _ = _from_event_dict({"source_id": "x", "url": "u", "url_hash": "h",
|
||||
"title": "t", "event": {
|
||||
"sentiment": "neutral",
|
||||
"importance": 1,
|
||||
"event_type": "其他",
|
||||
}})
|
||||
assert article.publish_time is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# remote 工具
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_chunked_splits_evenly() -> None:
|
||||
assert _chunked(list(range(7)), 3) == [[0, 1, 2], [3, 4, 5], [6]]
|
||||
assert _chunked([], 3) == []
|
||||
assert _chunked([1, 2, 3], 10) == [[1, 2, 3]]
|
||||
|
||||
|
||||
def test_dashscope_batch_limit_is_10() -> None:
|
||||
assert DASHSCOPE_BATCH_LIMIT == 10
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Provider type 解析
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_resolve_provider_type_dashscope_aliases(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("EMBEDDING_PROVIDER", raising=False)
|
||||
assert resolve_provider_type("dashscope") == EmbeddingProviderType.DASHSCOPE
|
||||
assert resolve_provider_type("qwen") == EmbeddingProviderType.DASHSCOPE
|
||||
assert resolve_provider_type("remote") == EmbeddingProviderType.DASHSCOPE
|
||||
|
||||
|
||||
def test_resolve_provider_type_local_aliases(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in ("local", "local-bge", "bge", "bge-m3"):
|
||||
assert resolve_provider_type(name) == EmbeddingProviderType.LOCAL_BGE
|
||||
|
||||
|
||||
def test_resolve_provider_type_default_is_dashscope(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("EMBEDDING_PROVIDER", raising=False)
|
||||
assert resolve_provider_type() == EmbeddingProviderType.DASHSCOPE
|
||||
|
||||
|
||||
def test_resolve_provider_type_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("EMBEDDING_PROVIDER", "local-bge")
|
||||
assert resolve_provider_type() == EmbeddingProviderType.LOCAL_BGE
|
||||
|
||||
|
||||
def test_resolve_provider_type_scene_override(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""configs/llm_models.yaml 的 scenes.embedding.provider 优先于 .env。"""
|
||||
monkeypatch.delenv("EMBEDDING_PROVIDER", raising=False)
|
||||
monkeypatch.setattr(
|
||||
"embedding.factory.load_scene_config",
|
||||
lambda scene: {"provider": "local-bge"} if scene == "embedding" else {},
|
||||
)
|
||||
assert resolve_provider_type() == EmbeddingProviderType.LOCAL_BGE
|
||||
|
||||
|
||||
def test_resolve_provider_type_explicit_arg_beats_scene(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""显式参数优先级最高,覆盖 YAML 场景。"""
|
||||
monkeypatch.setattr(
|
||||
"embedding.factory.load_scene_config",
|
||||
lambda scene: {"provider": "local-bge"} if scene == "embedding" else {},
|
||||
)
|
||||
assert resolve_provider_type("dashscope") == EmbeddingProviderType.DASHSCOPE
|
||||
|
||||
|
||||
def test_resolve_provider_type_unknown_raises() -> None:
|
||||
with pytest.raises(EmbeddingError):
|
||||
resolve_provider_type("anthropic-emb")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# DashScope 同步/异步(mock 网络)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_dashscope_sync_embed_batch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
provider = DashScopeEmbeddingProvider(model="text-embedding-v3")
|
||||
fake = MagicMock()
|
||||
fake.embeddings.create = MagicMock(
|
||||
side_effect=lambda model, input: _embedding_response([[0.1] * 1024] * len(input))
|
||||
)
|
||||
provider._client = fake
|
||||
out = provider.embed_batch(["a", "b", "c"])
|
||||
assert len(out) == 3
|
||||
assert all(len(v) == 1024 for v in out)
|
||||
|
||||
|
||||
def test_dashscope_sync_chunks_when_over_batch_limit(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
provider = DashScopeEmbeddingProvider()
|
||||
fake = MagicMock()
|
||||
fake.embeddings.create = MagicMock(
|
||||
side_effect=lambda model, input: _embedding_response([[0.0] * 1024] * len(input))
|
||||
)
|
||||
provider._client = fake
|
||||
texts = [f"t{i}" for i in range(25)] # > 10 -> 应分 3 批 (10+10+5)
|
||||
provider.embed_batch(texts)
|
||||
assert fake.embeddings.create.call_count == 3
|
||||
|
||||
|
||||
def test_dashscope_sync_retries_then_succeeds(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
provider = DashScopeEmbeddingProvider(max_attempts=3)
|
||||
fake = MagicMock()
|
||||
fake.embeddings.create = MagicMock(
|
||||
side_effect=[
|
||||
RuntimeError("rate-limit"),
|
||||
_embedding_response([[0.1] * 1024]),
|
||||
]
|
||||
)
|
||||
provider._client = fake
|
||||
out = provider.embed_batch(["x"])
|
||||
assert len(out) == 1
|
||||
assert fake.embeddings.create.call_count == 2
|
||||
|
||||
|
||||
def test_dashscope_sync_gives_up(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
provider = DashScopeEmbeddingProvider(max_attempts=2)
|
||||
fake = MagicMock()
|
||||
fake.embeddings.create = MagicMock(side_effect=RuntimeError("net"))
|
||||
provider._client = fake
|
||||
with pytest.raises(EmbeddingError) as exc:
|
||||
provider.embed_batch(["x"])
|
||||
assert exc.value.attempts == 2
|
||||
|
||||
|
||||
def test_dashscope_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False)
|
||||
with pytest.raises(EmbeddingError):
|
||||
DashScopeEmbeddingProvider()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dashscope_async_embed_batch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
provider = DashScopeAsyncEmbeddingProvider()
|
||||
fake = MagicMock()
|
||||
fake.embeddings.create = AsyncMock(
|
||||
side_effect=lambda model, input: _embedding_response([[0.1] * 1024] * len(input))
|
||||
)
|
||||
provider._client = fake
|
||||
out = await provider.embed_batch(["a", "b"])
|
||||
assert len(out) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dashscope_async_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
provider = DashScopeAsyncEmbeddingProvider()
|
||||
fake = MagicMock()
|
||||
fake.embeddings.create = AsyncMock(
|
||||
side_effect=lambda model, input: _embedding_response([[0.0] * 1024] * len(input))
|
||||
)
|
||||
provider._client = fake
|
||||
texts = [f"t{i}" for i in range(15)] # 2 批 (10+5)
|
||||
await provider.embed_batch(texts)
|
||||
assert fake.embeddings.create.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dashscope_async_retries(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
provider = DashScopeAsyncEmbeddingProvider(max_attempts=2)
|
||||
fake = MagicMock()
|
||||
fake.embeddings.create = AsyncMock(
|
||||
side_effect=[RuntimeError("transient"), _embedding_response([[0.0] * 1024])]
|
||||
)
|
||||
provider._client = fake
|
||||
out = await provider.embed_batch(["a"])
|
||||
assert len(out) == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# factory
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_make_sync_provider_dashscope(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
p = make_sync_provider("dashscope")
|
||||
assert isinstance(p, EmbeddingProvider)
|
||||
assert p.name == "dashscope"
|
||||
assert p.dim == 1024
|
||||
|
||||
|
||||
def test_make_async_provider_dashscope(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test")
|
||||
p = make_async_provider("qwen")
|
||||
assert isinstance(p, AsyncEmbeddingProvider)
|
||||
assert p.name == "dashscope"
|
||||
|
||||
|
||||
def test_make_sync_provider_local_without_st_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""无 sentence-transformers 时,本地 provider 应给出友好错误。"""
|
||||
import sys
|
||||
# 模拟 sentence_transformers 缺失
|
||||
monkeypatch.setitem(sys.modules, "sentence_transformers", None)
|
||||
with pytest.raises(EmbeddingError) as exc:
|
||||
make_sync_provider("local")
|
||||
assert "sentence-transformers" in str(exc.value)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# EmbeddingResult 模型 + 序列化
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_embedding_result_serializes(tmp_path: Path) -> None:
|
||||
r = EmbeddingResult(
|
||||
url_hash="abc",
|
||||
source_id="cls",
|
||||
title="t",
|
||||
text="text",
|
||||
vector=[0.1, 0.2, 0.3],
|
||||
dim=3,
|
||||
provider="dashscope",
|
||||
model="text-embedding-v3",
|
||||
)
|
||||
p = tmp_path / "r.json"
|
||||
p.write_text(r.model_dump_json(), encoding="utf-8")
|
||||
obj = json.loads(p.read_text(encoding="utf-8"))
|
||||
assert obj["dim"] == 3
|
||||
assert obj["vector"] == [0.1, 0.2, 0.3]
|
||||
|
||||
|
||||
def test_embedding_result_short_summary() -> None:
|
||||
r = EmbeddingResult(
|
||||
url_hash="abc",
|
||||
source_id="cls",
|
||||
title="宁德时代签约",
|
||||
text="x",
|
||||
vector=[0.0] * 4,
|
||||
dim=4,
|
||||
provider="dashscope",
|
||||
model="text-embedding-v3",
|
||||
)
|
||||
s = r.short_summary()
|
||||
assert "cls" in s and "dim=4" in s and "dashscope" in s
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 集成式: compose_text + 假异步 provider
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class _FakeAsyncProvider(AsyncEmbeddingProvider):
|
||||
name = "fake"
|
||||
model = "fake-1"
|
||||
dim = 8
|
||||
|
||||
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
|
||||
return [[float(len(t))] * self.dim for t in texts]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fake_async_provider_round_trip() -> None:
|
||||
art = _article()
|
||||
text = compose_text(art)
|
||||
async with _FakeAsyncProvider() as p:
|
||||
v = (await p.embed_batch([text]))[0]
|
||||
assert len(v) == 8
|
||||
assert v[0] == float(len(text))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fake_async_provider_concurrent_batches() -> None:
|
||||
p = _FakeAsyncProvider()
|
||||
res = await asyncio.gather(
|
||||
p.embed_batch(["a", "bb"]),
|
||||
p.embed_batch(["ccc"]),
|
||||
)
|
||||
assert res[0][0][0] == 1.0
|
||||
assert res[0][1][0] == 2.0
|
||||
assert res[1][0][0] == 3.0
|
||||
@@ -0,0 +1,481 @@
|
||||
"""M2 正文提取模块单元测试。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from extractor import Article, ExtractError, extract_article
|
||||
from extractor.parser import (
|
||||
_clean_content,
|
||||
_count_chinese,
|
||||
_fallback_time_from_html,
|
||||
_is_boilerplate,
|
||||
_is_reasonable_time,
|
||||
_normalize_time,
|
||||
_refine_title,
|
||||
_resolve_publish_time,
|
||||
_url_hash,
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 构造测试 HTML
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
SAMPLE_HTML = """
|
||||
<html>
|
||||
<head><title>宁德时代发布麒麟电池 - 财联社</title></head>
|
||||
<body>
|
||||
<header><nav>首页 财经 股票</nav></header>
|
||||
<div class="ad">广告:抢购618</div>
|
||||
|
||||
<article>
|
||||
<h1>宁德时代发布新一代麒麟电池 能量密度达 255Wh/kg</h1>
|
||||
<div class="meta">
|
||||
<span class="time">2026-06-15 14:30:00</span>
|
||||
<span class="author">记者 张三</span>
|
||||
</div>
|
||||
<p>财联社6月15日电,宁德时代今日正式发布了新一代麒麟电池产品,能量密度达到 255 Wh/kg,
|
||||
显著优于上一代产品的 213 Wh/kg。该产品定位高端电动车市场。</p>
|
||||
<p>据公司公告,该电池将于2026年第三季度量产,首批应用于多款新能源汽车。
|
||||
公司股价应声上涨5.2%,创年内新高。</p>
|
||||
<p>分析师认为,这将进一步巩固宁德时代在动力电池领域的全球领先地位,
|
||||
预计2026年公司动力电池出货量将同比增长30%以上。</p>
|
||||
<p>同时,公司还披露了海外建厂计划,德国与匈牙利工厂将于2027年投产。</p>
|
||||
</article>
|
||||
|
||||
<aside class="related">
|
||||
<h3>相关阅读</h3>
|
||||
<ul><li><a href="/a/1.html">比亚迪发布刀片电池升级版</a></li>
|
||||
<li><a href="/a/2.html">国轩高科上半年扭亏为盈</a></li></ul>
|
||||
</aside>
|
||||
|
||||
<section class="comments">
|
||||
<h3>评论 (88)</h3>
|
||||
<p>用户A: 太牛了!</p>
|
||||
<p>用户B: 行业要变天</p>
|
||||
</section>
|
||||
|
||||
<footer>版权所有 财联社</footer>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
|
||||
SHORT_HTML = """
|
||||
<html><body><h1>很短</h1><p>太短</p></body></html>
|
||||
"""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 主提取流程
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_extract_basic() -> None:
|
||||
article = extract_article(
|
||||
html=SAMPLE_HTML,
|
||||
source_id="cls",
|
||||
url="https://www.cls.cn/detail/123456",
|
||||
)
|
||||
assert isinstance(article, Article)
|
||||
assert article.source_id == "cls"
|
||||
assert article.source_name == "财联社"
|
||||
assert "宁德时代" in article.title
|
||||
assert "麒麟电池" in article.title
|
||||
assert article.url_hash == _url_hash("https://www.cls.cn/detail/123456")
|
||||
|
||||
# 关键正文短语必须保留
|
||||
assert "255 Wh/kg" in article.content or "255Wh/kg" in article.content
|
||||
assert "海外建厂" in article.content
|
||||
|
||||
# 时间被解析
|
||||
assert article.publish_time is not None
|
||||
assert article.publish_time.year == 2026
|
||||
assert article.publish_time.month == 6
|
||||
assert article.publish_time.day == 15
|
||||
|
||||
# 字数
|
||||
assert article.word_count > 30
|
||||
|
||||
|
||||
def test_extract_h1_preferred_over_short_title() -> None:
|
||||
"""<title> 仅有站名时应优先 H1。"""
|
||||
html = """
|
||||
<html><head><title>财联社</title></head><body>
|
||||
<h1>这是一个明显更长更详细的真实文章标题</h1>
|
||||
<article>
|
||||
<p>正文段落一,正文段落一,正文段落一,正文段落一,正文段落一。</p>
|
||||
<p>正文段落二,正文段落二,正文段落二,正文段落二,正文段落二。</p>
|
||||
<p>正文段落三,正文段落三,正文段落三,正文段落三,正文段落三。</p>
|
||||
</article></body></html>
|
||||
"""
|
||||
article = extract_article(html, "cls", "https://www.cls.cn/detail/1")
|
||||
assert article.title == "这是一个明显更长更详细的真实文章标题"
|
||||
|
||||
|
||||
def test_extract_too_short_raises() -> None:
|
||||
with pytest.raises(ExtractError):
|
||||
extract_article(SHORT_HTML, "cls", "https://www.cls.cn/detail/2")
|
||||
|
||||
|
||||
def test_extract_empty_html_raises() -> None:
|
||||
with pytest.raises(ExtractError):
|
||||
extract_article("", "cls", "https://www.cls.cn/detail/3")
|
||||
with pytest.raises(ExtractError):
|
||||
extract_article(" \n\n ", "cls", "https://www.cls.cn/detail/4")
|
||||
|
||||
|
||||
def test_extract_unknown_source_has_no_source_name() -> None:
|
||||
article = extract_article(SAMPLE_HTML, "unknown_src", "https://example.com/a/1")
|
||||
assert article.source_name is None
|
||||
assert article.source_id == "unknown_src"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 模板兜底检测(M2.2)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_is_boilerplate_eastmoney_legal_disclaimer() -> None:
|
||||
text = (
|
||||
"郑重声明: 1.根据《证券法》规定,禁止编造、传播虚假信息或者误导性信息,"
|
||||
"扰乱证券市场;2.用户在本社区发表的所有资料、言论等仅代表个人观点。"
|
||||
)
|
||||
is_bp, reason = _is_boilerplate(text)
|
||||
assert is_bp
|
||||
assert reason is not None and "郑重声明" in reason
|
||||
|
||||
|
||||
def test_is_boilerplate_yicai_ad_copyright() -> None:
|
||||
text = (
|
||||
"第一财经广告合作,请点击这里。此内容为第一财经原创,著作权归第一财经所有。"
|
||||
"未经第一财经书面授权,不得以任何方式加以使用。"
|
||||
)
|
||||
is_bp, reason = _is_boilerplate(text)
|
||||
assert is_bp
|
||||
|
||||
|
||||
def test_is_boilerplate_sina_embedded_post() -> None:
|
||||
text = (
|
||||
"北京红竹 今天 16:02:49 本质上就是一句话:资金开始从核心抱团,进入扩散阶段。"
|
||||
"银行这边今天已经明显走弱。银行和科技之间还是跷跷板关系。" * 2
|
||||
)
|
||||
is_bp, _ = _is_boilerplate(text)
|
||||
assert is_bp
|
||||
|
||||
|
||||
def test_is_boilerplate_long_real_article_passes() -> None:
|
||||
"""真实长篇文章中即使提到"郑重声明"等词,长度 > 800 不应误判。"""
|
||||
legitimate = (
|
||||
"公司发布郑重声明,回应近期市场关注的多项议题。根据《证券法》及相关法律法规要求,"
|
||||
"公司将严格履行信息披露义务。" * 30 # ~1200 字
|
||||
)
|
||||
assert len(legitimate) > 800
|
||||
is_bp, _ = _is_boilerplate(legitimate)
|
||||
assert not is_bp
|
||||
|
||||
|
||||
def test_is_boilerplate_empty_returns_false() -> None:
|
||||
assert _is_boilerplate("") == (False, None)
|
||||
assert _is_boilerplate("普通正文,没有任何模板特征,长度也合理。" * 5) == (False, None)
|
||||
|
||||
|
||||
def test_extract_article_raises_on_boilerplate() -> None:
|
||||
"""模拟一篇 GNE 兜底失败的页面,extract_article 应抛 ExtractError。"""
|
||||
html = """
|
||||
<html><head><title>东方财富</title></head><body>
|
||||
<h1>是否有中国船只通过霍尔木兹海峡？外交部回应</h1>
|
||||
<article>
|
||||
郑重声明: 1.根据《证券法》规定,禁止编造、传播虚假信息或者误导性信息,
|
||||
扰乱证券市场;2.用户在本社区发表的所有资料、言论等仅代表个人观点,与本网站立场无关。
|
||||
《东方财富社区管理规定》
|
||||
</article>
|
||||
</body></html>
|
||||
"""
|
||||
with pytest.raises(ExtractError) as exc:
|
||||
extract_article(
|
||||
html,
|
||||
source_id="eastmoney",
|
||||
url="https://finance.eastmoney.com/a/123.html",
|
||||
)
|
||||
assert "模板" in str(exc.value)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# title 校正
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_refine_title_uses_h1_when_gne_is_substring() -> None:
|
||||
html = "<html><body><h1>完整真实标题</h1></body></html>"
|
||||
assert _refine_title(html, "完整真实") == "完整真实标题"
|
||||
|
||||
|
||||
def test_refine_title_uses_h1_when_gne_has_site_suffix() -> None:
|
||||
html = "<html><body><h1>真实标题</h1></body></html>"
|
||||
assert _refine_title(html, "真实标题 - 财联社") == "真实标题"
|
||||
|
||||
|
||||
def test_refine_title_falls_back_to_gne_when_no_h1() -> None:
|
||||
html = "<html><body><p>x</p></body></html>"
|
||||
assert _refine_title(html, "GNE 标题") == "GNE 标题"
|
||||
|
||||
|
||||
def test_refine_title_returns_empty_when_both_missing() -> None:
|
||||
assert _refine_title("<html></html>", "") == ""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# content 清理
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_clean_content_strips_repeated_header() -> None:
|
||||
raw = "宁德时代发布新一代麒麟电池\n2026-06-15 14:30:00\n记者 张三\n\n正文第一段。\n\n\n\n正文第二段。"
|
||||
cleaned = _clean_content(
|
||||
raw,
|
||||
title="宁德时代发布新一代麒麟电池",
|
||||
author="记者 张三",
|
||||
time_raw="2026-06-15 14:30:00",
|
||||
)
|
||||
assert cleaned.startswith("正文第一段")
|
||||
assert "正文第二段" in cleaned
|
||||
# 连续 4 个 \n 应被压缩为 \n\n
|
||||
assert "\n\n\n" not in cleaned
|
||||
|
||||
|
||||
def test_clean_content_keeps_body_when_no_header_match() -> None:
|
||||
raw = "段落一\n段落二\n段落三"
|
||||
cleaned = _clean_content(raw, title="完全不同的标题", author=None, time_raw=None)
|
||||
assert cleaned == "段落一\n段落二\n段落三"
|
||||
|
||||
|
||||
def test_clean_content_handles_empty() -> None:
|
||||
assert _clean_content("", "", None, None) == ""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 时间标准化
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_normalize_time_iso() -> None:
|
||||
dt = _normalize_time("2026-06-15 14:30:00")
|
||||
assert dt is not None and dt == datetime(2026, 6, 15, 14, 30, 0)
|
||||
|
||||
|
||||
def test_normalize_time_chinese_format() -> None:
|
||||
dt = _normalize_time("2026年6月15日 14时30分")
|
||||
assert dt is not None
|
||||
assert dt.year == 2026 and dt.month == 6 and dt.day == 15
|
||||
assert dt.hour == 14 and dt.minute == 30
|
||||
|
||||
|
||||
def test_normalize_time_chinese_seconds() -> None:
|
||||
dt = _normalize_time("2026年1月1日 9时5分3秒")
|
||||
assert dt is not None and dt.second == 3
|
||||
|
||||
|
||||
def test_normalize_time_invalid_returns_none() -> None:
|
||||
assert _normalize_time("不是时间") is None
|
||||
assert _normalize_time("") is None
|
||||
assert _normalize_time(None) is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 时间合理性 + HTML 兜底
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_is_reasonable_time_within_range() -> None:
|
||||
ref = datetime(2026, 6, 16, 12, 0, 0)
|
||||
assert _is_reasonable_time(datetime(2026, 6, 15, 8, 0), ref)
|
||||
assert _is_reasonable_time(datetime(2025, 12, 1, 0, 0), ref)
|
||||
# 同一天即将到来的时间
|
||||
assert _is_reasonable_time(datetime(2026, 6, 16, 23, 0), ref)
|
||||
|
||||
|
||||
def test_is_reasonable_time_rejects_far_future() -> None:
|
||||
ref = datetime(2026, 6, 16, 12, 0)
|
||||
assert not _is_reasonable_time(datetime(2026, 6, 18, 0, 0), ref)
|
||||
|
||||
|
||||
def test_is_reasonable_time_rejects_far_past() -> None:
|
||||
ref = datetime(2026, 6, 16, 12, 0)
|
||||
# 距 ref 超过 365 天 -> 不合理(模拟 GNE 抓到的 2019 年页脚时间)
|
||||
assert not _is_reasonable_time(datetime(2019, 1, 16, 10, 40), ref)
|
||||
|
||||
|
||||
def test_is_reasonable_time_strips_tzinfo() -> None:
|
||||
"""带时区的时间也能与 naive ref 比较。"""
|
||||
from datetime import UTC
|
||||
|
||||
aware = datetime(2026, 6, 16, 12, 0, tzinfo=UTC)
|
||||
assert _is_reasonable_time(aware, datetime(2026, 6, 16, 12, 0))
|
||||
|
||||
|
||||
def test_is_reasonable_time_handles_none() -> None:
|
||||
assert _is_reasonable_time(None) is False
|
||||
|
||||
|
||||
def test_fallback_time_from_html_finds_eastmoney_pattern() -> None:
|
||||
"""模拟 eastmoney 真实结构,兜底应能命中 .infos 内的中文日期。"""
|
||||
html = """
|
||||
<div class="infos">
|
||||
<div class="item">2026年06月16日 17:54</div>
|
||||
<div class="item">来源:发改委网站</div>
|
||||
</div>
|
||||
"""
|
||||
dt, raw = _fallback_time_from_html(html)
|
||||
assert dt is not None
|
||||
assert dt.year == 2026 and dt.month == 6 and dt.day == 16
|
||||
assert dt.hour == 17 and dt.minute == 54
|
||||
assert raw is not None and "2026" in raw
|
||||
|
||||
|
||||
def test_fallback_time_skips_unreasonable_dates() -> None:
|
||||
"""页脚备案/版权时间(2019-01-16)出现在前,真实发布时间在后,应跳过前者。"""
|
||||
html = """
|
||||
<html><body>
|
||||
<header>
|
||||
<div class="biaobei">备案号 京ICP-XXX 备案日期: 2019-01-16</div>
|
||||
</header>
|
||||
<div class="infos">
|
||||
<div class="item">2026年06月16日 17:54</div>
|
||||
</div>
|
||||
</body></html>
|
||||
"""
|
||||
dt, raw = _fallback_time_from_html(html)
|
||||
assert dt is not None
|
||||
assert dt.year == 2026
|
||||
assert "2026" in (raw or "")
|
||||
|
||||
|
||||
def test_fallback_time_returns_none_when_no_match() -> None:
|
||||
dt, raw = _fallback_time_from_html("<html><body>没有日期</body></html>")
|
||||
assert dt is None and raw is None
|
||||
|
||||
|
||||
def test_resolve_publish_time_uses_gne_when_reasonable() -> None:
|
||||
today = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
dt, raw = _resolve_publish_time("<html></html>", today)
|
||||
assert dt is not None
|
||||
assert raw == today
|
||||
|
||||
|
||||
def test_resolve_publish_time_falls_back_when_gne_unreasonable() -> None:
|
||||
"""模拟 eastmoney 场景:GNE 给出 2019-01-16,HTML 中含真实时间。"""
|
||||
html = '<div class="infos"><div class="item">2026年06月16日 17:54</div></div>'
|
||||
dt, raw = _resolve_publish_time(html, "2019-01-16 10:40:21")
|
||||
assert dt is not None and dt.year == 2026 and dt.month == 6 and dt.day == 16
|
||||
# raw 应反映兜底来源,而非原始 GNE 字符串
|
||||
assert "2026" in (raw or "")
|
||||
|
||||
|
||||
def test_resolve_publish_time_keeps_gne_raw_when_all_fail() -> None:
|
||||
"""GNE 不合理且 HTML 也无可用时间 -> publish_time None,raw 保留 GNE。"""
|
||||
dt, raw = _resolve_publish_time("<html><body>无</body></html>", "2010-01-01 00:00:00")
|
||||
assert dt is None
|
||||
assert raw == "2010-01-01 00:00:00"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 端到端:真实 eastmoney 结构应解析出 2026 时间
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_extract_article_uses_html_fallback_for_eastmoney_like() -> None:
|
||||
"""模拟真实 eastmoney 文章页:GNE 拿到页脚错误时间,extractor 应兜底。"""
|
||||
html = """
|
||||
<html><head><title>事关六张网建设</title></head><body>
|
||||
<footer>备案信息发布时间 2019-01-16 10:40:21</footer>
|
||||
<div id="topbox" class="topbox">
|
||||
<div class="title">事关“六张网”建设 国家发展改革委召开重要座谈会</div>
|
||||
<div class="tipbox">
|
||||
<div class="infos">
|
||||
<div class="item">2026年06月16日 17:54</div>
|
||||
<div class="item">来源:发改委网站</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<article>
|
||||
<h1>事关“六张网”建设 国家发展改革委召开重要座谈会</h1>
|
||||
<p>近日,国家发展改革委召开座谈会,围绕加快推进“六张网”建设的具体举措进行专题研讨。</p>
|
||||
<p>会议指出,“六张网”建设关系国家长远发展,要从体制机制、关键技术、重点项目三方面协同推进。</p>
|
||||
<p>与会专家就资金保障、跨部门协调、技术标准统一等议题展开了深入交流。</p>
|
||||
</article>
|
||||
</body></html>
|
||||
"""
|
||||
article = extract_article(
|
||||
html,
|
||||
source_id="eastmoney",
|
||||
url="https://finance.eastmoney.com/a/123456789.html",
|
||||
)
|
||||
assert article.publish_time is not None
|
||||
assert article.publish_time.year == 2026
|
||||
assert article.publish_time.month == 6
|
||||
assert article.publish_time.day == 16
|
||||
assert "2026" in (article.publish_time_raw or "")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 工具
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_url_hash_stable_and_short() -> None:
|
||||
h1 = _url_hash("https://example.com/a")
|
||||
h2 = _url_hash("https://example.com/a")
|
||||
h3 = _url_hash("https://example.com/b")
|
||||
assert h1 == h2 != h3
|
||||
assert len(h1) == 16
|
||||
|
||||
|
||||
def test_count_chinese() -> None:
|
||||
assert _count_chinese("hello 你好 world") == 2
|
||||
assert _count_chinese("ABC123") == 0
|
||||
assert _count_chinese("中文测试") == 4
|
||||
|
||||
|
||||
def test_article_short_summary() -> None:
|
||||
a = Article(
|
||||
source_id="cls",
|
||||
url="https://x/1",
|
||||
url_hash="abc123",
|
||||
title="测试标题",
|
||||
content="一段正文 " * 20,
|
||||
publish_time=datetime(2026, 6, 15, 14, 30),
|
||||
word_count=20,
|
||||
)
|
||||
s = a.short_summary()
|
||||
assert "cls" in s
|
||||
assert "2026-06-15" in s
|
||||
assert "测试标题" in s
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Article 模型字段约束
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_article_requires_min_length_title_content() -> None:
|
||||
with pytest.raises(Exception): # noqa: B017 - pydantic ValidationError
|
||||
Article(
|
||||
source_id="cls",
|
||||
url="https://x/1",
|
||||
url_hash="h",
|
||||
title="",
|
||||
content="",
|
||||
)
|
||||
|
||||
|
||||
def test_article_serializes_to_json(tmp_path: Path) -> None:
|
||||
a = Article(
|
||||
source_id="cls",
|
||||
url="https://x/1",
|
||||
url_hash="abc",
|
||||
title="标题",
|
||||
content="一二三四五六七八九十一二三四五六七八九十一二三四五六七八九十",
|
||||
publish_time=datetime(2026, 6, 15, 14, 30),
|
||||
word_count=30,
|
||||
)
|
||||
p = tmp_path / "a.json"
|
||||
p.write_text(a.model_dump_json(), encoding="utf-8")
|
||||
obj = json.loads(p.read_text(encoding="utf-8"))
|
||||
assert obj["source_id"] == "cls"
|
||||
assert obj["publish_time"].startswith("2026-06-15")
|
||||
@@ -0,0 +1,268 @@
|
||||
"""增量处理与 pipeline 断点续跑测试。
|
||||
|
||||
覆盖:
|
||||
- M2/M4/M5 脚本的「产物存在即跳过」过滤逻辑
|
||||
- scheduler.pipeline 断点状态记录与 --resume 续跑逻辑
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
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
|
||||
@@ -0,0 +1,500 @@
|
||||
"""M4 LLM 投资事件抽取测试。
|
||||
|
||||
不依赖真实 LLM API,所有调用通过 mock 注入响应。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from extractor import Article
|
||||
from llm import (
|
||||
EVENT_TYPES,
|
||||
EventExtraction,
|
||||
ExtractedEvent,
|
||||
LLMCallError,
|
||||
PromptTemplate,
|
||||
Sentiment,
|
||||
extract_event,
|
||||
extract_event_async,
|
||||
load_llm_config,
|
||||
parse_event_json,
|
||||
)
|
||||
from llm.client import LLMConfig
|
||||
from llm.extractor import _extract_json_object
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# fixtures
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _article(
|
||||
*,
|
||||
title: str = "宁德时代签订 100GWh 长期供货协议",
|
||||
content: str = "宁德时代(300750)与某车企签 5 年 100GWh 协议,涉及金额超 1500 亿。" * 3,
|
||||
publish_time: datetime | None = datetime(2026, 6, 16, 10, 0),
|
||||
) -> Article:
|
||||
return Article(
|
||||
source_id="cls",
|
||||
url="https://www.cls.cn/detail/1",
|
||||
url_hash="abc1234567890000",
|
||||
title=title,
|
||||
content=content,
|
||||
publish_time=publish_time,
|
||||
word_count=len(content),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_config() -> LLMConfig:
|
||||
return LLMConfig(
|
||||
provider="deepseek",
|
||||
model="deepseek-chat",
|
||||
api_key="sk-fake",
|
||||
base_url="https://api.deepseek.com",
|
||||
)
|
||||
|
||||
|
||||
def _mock_completion(content: str, prompt_tokens: int = 100, completion_tokens: int = 50) -> MagicMock:
|
||||
"""构造与 openai SDK 返回兼容的 mock 对象。"""
|
||||
msg = MagicMock()
|
||||
msg.content = content
|
||||
choice = MagicMock()
|
||||
choice.message = msg
|
||||
usage = MagicMock()
|
||||
usage.prompt_tokens = prompt_tokens
|
||||
usage.completion_tokens = completion_tokens
|
||||
resp = MagicMock()
|
||||
resp.choices = [choice]
|
||||
resp.usage = usage
|
||||
return resp
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# EventExtraction 模型校验
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_event_extraction_minimal_valid() -> None:
|
||||
e = EventExtraction(sentiment="positive", importance=4, event_type="重大合同")
|
||||
assert e.sentiment == Sentiment.POSITIVE
|
||||
assert e.importance == 4
|
||||
|
||||
|
||||
def test_event_extraction_normalizes_stock_codes() -> None:
|
||||
e = EventExtraction(
|
||||
stock_codes=["300750.sz", " 300750.SZ ", "abc", "12345", "600519"],
|
||||
sentiment="positive", importance=3, event_type="其他",
|
||||
)
|
||||
# 大小写 / 空白被规范;非法被过滤;去重
|
||||
assert e.stock_codes == ["300750.SZ", "600519"]
|
||||
|
||||
|
||||
def test_event_extraction_filters_empty_lists() -> None:
|
||||
e = EventExtraction(
|
||||
stock_codes=[], company_names=["", " ", "宁德时代", "宁德时代"],
|
||||
industries=[],
|
||||
sentiment="neutral", importance=1, event_type="其他",
|
||||
)
|
||||
assert e.company_names == ["宁德时代"]
|
||||
assert e.stock_codes == []
|
||||
|
||||
|
||||
def test_event_extraction_rejects_importance_out_of_range() -> None:
|
||||
with pytest.raises(Exception): # noqa: B017
|
||||
EventExtraction(sentiment="positive", importance=0, event_type="其他")
|
||||
with pytest.raises(Exception): # noqa: B017
|
||||
EventExtraction(sentiment="positive", importance=6, event_type="其他")
|
||||
|
||||
|
||||
def test_event_extraction_normalizes_blank_event_type() -> None:
|
||||
e = EventExtraction(sentiment="neutral", importance=1, event_type=" ")
|
||||
assert e.event_type == "其他"
|
||||
|
||||
|
||||
def test_event_types_constant_includes_common() -> None:
|
||||
for must in ["业绩预告", "合作签约", "监管处罚", "其他"]:
|
||||
assert must in EVENT_TYPES
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# ExtractedEvent.sources 多源字段
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_extracted_event_sources_defaults_to_main_source() -> None:
|
||||
"""未提供 sources 时兜底为 [source_id](兼容旧产物)。"""
|
||||
ev = ExtractedEvent(
|
||||
source_id="cls", url="https://x/1", url_hash="h1", title="t",
|
||||
event=EventExtraction(sentiment="positive", importance=3, event_type="重大合同"),
|
||||
provider="deepseek", model="m",
|
||||
)
|
||||
assert ev.sources == ["cls"]
|
||||
|
||||
|
||||
def test_extracted_event_sources_keeps_main_first_and_dedup() -> None:
|
||||
"""sources 保主源居首、去重保序。"""
|
||||
ev = ExtractedEvent(
|
||||
source_id="cls", url="https://x/1", url_hash="h1", title="t",
|
||||
sources=["sina", "cls", "eastmoney", "sina"],
|
||||
event=EventExtraction(sentiment="neutral", importance=2, event_type="其他"),
|
||||
provider="deepseek", model="m",
|
||||
)
|
||||
assert ev.sources == ["cls", "sina", "eastmoney"]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# JSON 提取与解析
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_extract_json_object_strips_fence() -> None:
|
||||
s = '```json\n{"a": 1}\n```'
|
||||
assert _extract_json_object(s) == '{"a": 1}'
|
||||
|
||||
|
||||
def test_extract_json_object_picks_first_object() -> None:
|
||||
s = '前置说明\n{"a": 1}\n更多文字'
|
||||
assert _extract_json_object(s) == '{"a": 1}'
|
||||
|
||||
|
||||
def test_extract_json_object_handles_nested() -> None:
|
||||
s = '{"a": {"b": 2}}'
|
||||
assert _extract_json_object(s) == '{"a": {"b": 2}}'
|
||||
|
||||
|
||||
def test_parse_event_json_ok() -> None:
|
||||
raw = json.dumps({
|
||||
"stock_codes": ["300750.SZ"],
|
||||
"company_names": ["宁德时代"],
|
||||
"industries": ["动力电池"],
|
||||
"sentiment": "positive",
|
||||
"importance": 5,
|
||||
"event_type": "重大合同",
|
||||
"summary": "签订长期供货协议",
|
||||
})
|
||||
e = parse_event_json(raw)
|
||||
assert e.sentiment == Sentiment.POSITIVE
|
||||
assert e.stock_codes == ["300750.SZ"]
|
||||
|
||||
|
||||
def test_parse_event_json_invalid_json_raises() -> None:
|
||||
with pytest.raises(LLMCallError):
|
||||
parse_event_json("not a json")
|
||||
|
||||
|
||||
def test_parse_event_json_non_object_raises() -> None:
|
||||
with pytest.raises(LLMCallError):
|
||||
parse_event_json('["array"]')
|
||||
|
||||
|
||||
def test_parse_event_json_schema_invalid_raises() -> None:
|
||||
with pytest.raises(LLMCallError):
|
||||
parse_event_json('{"importance": 99}') # 缺 sentiment + event_type 且 importance 越界
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# PromptTemplate
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_prompt_template_renders_placeholders(tmp_path: Path) -> None:
|
||||
tpl_file = tmp_path / "tpl.md"
|
||||
tpl_file.write_text(
|
||||
"标题:{title}\n时间:{publish_time}\n源:{source_name}\n内容:\n{content}\nEND",
|
||||
encoding="utf-8",
|
||||
)
|
||||
tpl = PromptTemplate(tpl_file)
|
||||
art = _article()
|
||||
rendered = tpl.render(art)
|
||||
assert "标题:" + art.title in rendered
|
||||
assert "2026-06-16" in rendered
|
||||
assert "源:财联社" in rendered or "源:cls" in rendered # source_name 默认空,落到 source_id
|
||||
assert art.content[:30] in rendered
|
||||
|
||||
|
||||
def test_prompt_template_truncates_long_content(tmp_path: Path) -> None:
|
||||
tpl_file = tmp_path / "tpl.md"
|
||||
tpl_file.write_text("{content}", encoding="utf-8")
|
||||
tpl = PromptTemplate(tpl_file)
|
||||
art = _article(content="字" * 20000)
|
||||
rendered = tpl.render(art)
|
||||
assert "[正文过长已截断]" in rendered
|
||||
assert len(rendered) < 20000
|
||||
|
||||
|
||||
def test_prompt_template_default_path_loads() -> None:
|
||||
"""项目内置 prompts/event_extraction.md 必须可加载,作为回归保护。"""
|
||||
real = Path("prompts/event_extraction.md")
|
||||
if not real.is_file():
|
||||
pytest.skip("prompts/event_extraction.md 未找到")
|
||||
tpl = PromptTemplate(real)
|
||||
out = tpl.render(_article())
|
||||
assert "{title}" not in out
|
||||
assert "{content}" not in out
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# extract_event(同步,带重试)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_extract_event_succeeds_first_try(fake_config: LLMConfig) -> None:
|
||||
client = MagicMock()
|
||||
raw = json.dumps({
|
||||
"stock_codes": ["300750.SZ"],
|
||||
"company_names": ["宁德时代"],
|
||||
"industries": ["动力电池"],
|
||||
"sentiment": "positive",
|
||||
"importance": 5,
|
||||
"event_type": "重大合同",
|
||||
"summary": "100GWh 合作",
|
||||
})
|
||||
client.chat.completions.create = MagicMock(return_value=_mock_completion(raw))
|
||||
|
||||
result = extract_event(client, fake_config, _article())
|
||||
assert isinstance(result, ExtractedEvent)
|
||||
assert result.attempts == 1
|
||||
assert result.event.sentiment == Sentiment.POSITIVE
|
||||
assert result.provider == "deepseek"
|
||||
assert result.prompt_tokens == 100
|
||||
|
||||
|
||||
def test_extract_event_retries_on_invalid_json(fake_config: LLMConfig) -> None:
|
||||
"""第 1/2 次返回非法 JSON,第 3 次成功。"""
|
||||
valid = json.dumps({
|
||||
"sentiment": "neutral", "importance": 1, "event_type": "其他",
|
||||
})
|
||||
client = MagicMock()
|
||||
client.chat.completions.create = MagicMock(side_effect=[
|
||||
_mock_completion("not a json"),
|
||||
_mock_completion('{"sentiment":"???"}'), # schema 校验失败
|
||||
_mock_completion(valid),
|
||||
])
|
||||
result = extract_event(client, fake_config, _article(), max_attempts=3)
|
||||
assert result.attempts == 3
|
||||
assert client.chat.completions.create.call_count == 3
|
||||
|
||||
|
||||
def test_extract_event_gives_up_after_max(fake_config: LLMConfig) -> None:
|
||||
client = MagicMock()
|
||||
client.chat.completions.create = MagicMock(
|
||||
return_value=_mock_completion("not a json")
|
||||
)
|
||||
with pytest.raises(LLMCallError) as exc:
|
||||
extract_event(client, fake_config, _article(), max_attempts=2)
|
||||
assert exc.value.attempts == 2
|
||||
assert client.chat.completions.create.call_count == 2
|
||||
|
||||
|
||||
def test_extract_event_handles_network_exception(fake_config: LLMConfig) -> None:
|
||||
client = MagicMock()
|
||||
valid = json.dumps({
|
||||
"sentiment": "negative", "importance": 3, "event_type": "监管处罚",
|
||||
})
|
||||
client.chat.completions.create = MagicMock(side_effect=[
|
||||
TimeoutError("net hang"),
|
||||
_mock_completion(valid),
|
||||
])
|
||||
result = extract_event(client, fake_config, _article(), max_attempts=2)
|
||||
assert result.attempts == 2
|
||||
assert result.event.sentiment == Sentiment.NEGATIVE
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# extract_event_async
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_event_async_succeeds(fake_config: LLMConfig) -> None:
|
||||
client = MagicMock()
|
||||
valid = json.dumps({
|
||||
"stock_codes": ["600519"],
|
||||
"company_names": ["贵州茅台"],
|
||||
"sentiment": "neutral",
|
||||
"importance": 2,
|
||||
"event_type": "财报披露",
|
||||
"summary": "披露半年报",
|
||||
})
|
||||
client.chat.completions.create = AsyncMock(return_value=_mock_completion(valid))
|
||||
|
||||
sem = asyncio.Semaphore(2)
|
||||
result = await extract_event_async(
|
||||
client, fake_config, _article(), semaphore=sem,
|
||||
)
|
||||
assert result.event.stock_codes == ["600519"]
|
||||
assert result.event.event_type == "财报披露"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_event_async_retries(fake_config: LLMConfig) -> None:
|
||||
valid = json.dumps({"sentiment": "positive", "importance": 4, "event_type": "其他"})
|
||||
client = MagicMock()
|
||||
client.chat.completions.create = AsyncMock(side_effect=[
|
||||
ValueError("transient"),
|
||||
_mock_completion(valid),
|
||||
])
|
||||
result = await extract_event_async(
|
||||
client, fake_config, _article(), max_attempts=2,
|
||||
)
|
||||
assert result.attempts == 2
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# load_llm_config
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_load_llm_config_deepseek_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LLM_PROVIDER", "deepseek")
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-test-deepseek")
|
||||
monkeypatch.setenv("DEEPSEEK_MODEL", "deepseek-v4-flash")
|
||||
monkeypatch.delenv("LLM_MODEL", raising=False)
|
||||
cfg = load_llm_config()
|
||||
assert cfg.provider == "deepseek"
|
||||
assert cfg.api_key == "sk-test-deepseek"
|
||||
assert cfg.model == "deepseek-v4-flash"
|
||||
assert "deepseek" in cfg.base_url
|
||||
|
||||
|
||||
def test_load_llm_config_qwen_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LLM_PROVIDER", "qwen")
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-test-qwen")
|
||||
monkeypatch.setenv("QWEN_MODEL", "qwen-plus")
|
||||
monkeypatch.delenv("LLM_MODEL", raising=False)
|
||||
cfg = load_llm_config()
|
||||
assert cfg.provider == "qwen"
|
||||
assert cfg.api_key == "sk-test-qwen"
|
||||
assert cfg.model == "qwen-plus"
|
||||
assert "dashscope" in cfg.base_url or "aliyuncs" in cfg.base_url
|
||||
|
||||
|
||||
def test_load_llm_config_unknown_provider_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
with pytest.raises(ValueError):
|
||||
load_llm_config(provider="anthropic")
|
||||
|
||||
|
||||
def test_load_llm_config_missing_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False)
|
||||
monkeypatch.setenv("DEEPSEEK_MODEL", "deepseek-v4-flash")
|
||||
with pytest.raises(ValueError, match="API key"):
|
||||
load_llm_config(provider="deepseek")
|
||||
|
||||
|
||||
def test_load_llm_config_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""去掉内置默认模型后：未显式配置模型必须报错（不再回退 deepseek-chat）。"""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-test")
|
||||
monkeypatch.delenv("DEEPSEEK_MODEL", raising=False)
|
||||
monkeypatch.delenv("LLM_MODEL", raising=False)
|
||||
with pytest.raises(ValueError, match="模型"):
|
||||
load_llm_config(provider="deepseek")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# load_llm_config —— configs/llm_models.yaml 场景配置
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _patch_scene(monkeypatch: pytest.MonkeyPatch, cfg: dict) -> None:
|
||||
"""替换场景加载,模拟 configs/llm_models.yaml 中的某场景配置。"""
|
||||
monkeypatch.setattr(
|
||||
"llm.client.load_scene_config",
|
||||
lambda scene: cfg if scene == "daily_report" else {},
|
||||
)
|
||||
|
||||
|
||||
def test_load_llm_config_scene_overrides_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""YAML 场景配置优先于 .env:provider / model / temperature / timeout / max_attempts。"""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-env")
|
||||
monkeypatch.setenv("DEEPSEEK_MODEL", "deepseek-env-model")
|
||||
monkeypatch.setenv("QWEN_API_KEY", "sk-qwen")
|
||||
_patch_scene(monkeypatch, {
|
||||
"provider": "qwen",
|
||||
"model": "qwen-max",
|
||||
"temperature": 0.5,
|
||||
"timeout_sec": 99,
|
||||
"max_attempts": 5,
|
||||
})
|
||||
cfg = load_llm_config(scene="daily_report")
|
||||
assert cfg.provider == "qwen"
|
||||
assert cfg.model == "qwen-max"
|
||||
assert cfg.api_key == "sk-qwen"
|
||||
assert cfg.temperature == 0.5
|
||||
assert cfg.timeout_sec == 99
|
||||
assert cfg.max_attempts == 5
|
||||
|
||||
|
||||
def test_load_llm_config_scene_blank_fields_fall_back_to_env(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""YAML 场景未配置的字段(如 model 留空)回退 .env,保持向后兼容。"""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-env")
|
||||
monkeypatch.setenv("DEEPSEEK_MODEL", "deepseek-env-model")
|
||||
_patch_scene(monkeypatch, {"provider": "deepseek", "model": "", "temperature": 0.7})
|
||||
cfg = load_llm_config(scene="daily_report")
|
||||
assert cfg.provider == "deepseek"
|
||||
assert cfg.model == "deepseek-env-model"
|
||||
assert cfg.api_key == "sk-env"
|
||||
assert cfg.temperature == 0.7
|
||||
|
||||
|
||||
def test_load_llm_config_scene_api_key_env_name(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""api_key_env 指向自定义环境变量时,优先使用该变量。"""
|
||||
monkeypatch.setenv("MY_CUSTOM_KEY", "sk-custom")
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-default")
|
||||
monkeypatch.setenv("DEEPSEEK_MODEL", "deepseek-m")
|
||||
_patch_scene(monkeypatch, {
|
||||
"provider": "deepseek",
|
||||
"model": "deepseek-scene-m",
|
||||
"api_key_env": "MY_CUSTOM_KEY",
|
||||
})
|
||||
cfg = load_llm_config(scene="daily_report")
|
||||
assert cfg.api_key == "sk-custom"
|
||||
assert cfg.model == "deepseek-scene-m"
|
||||
|
||||
|
||||
def test_load_llm_config_scene_explicit_args_win(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""CLI/显式参数优先级最高,覆盖 YAML 场景。"""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-env")
|
||||
_patch_scene(monkeypatch, {"provider": "qwen", "model": "qwen-max"})
|
||||
monkeypatch.setenv("QWEN_API_KEY", "sk-qwen")
|
||||
cfg = load_llm_config(provider="deepseek", model="deepseek-chat", scene="daily_report")
|
||||
assert cfg.provider == "deepseek"
|
||||
assert cfg.model == "deepseek-chat"
|
||||
|
||||
|
||||
def test_load_llm_config_scene_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""场景与 .env 都未配置模型时必须报错(无内置兜底)。"""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-test")
|
||||
monkeypatch.delenv("DEEPSEEK_MODEL", raising=False)
|
||||
monkeypatch.delenv("LLM_MODEL", raising=False)
|
||||
_patch_scene(monkeypatch, {"provider": "deepseek", "model": ""})
|
||||
with pytest.raises(ValueError, match="模型"):
|
||||
load_llm_config(scene="daily_report")
|
||||
|
||||
|
||||
def test_load_llm_config_real_yaml_parseable() -> None:
|
||||
"""真实 configs/llm_models.yaml 必须可解析且包含全部场景(回归保护)。"""
|
||||
from configs.loader import load_defaults, load_scene_config
|
||||
|
||||
for scene in ("event_extraction", "daily_report", "stock_report", "embedding"):
|
||||
assert isinstance(load_scene_config(scene), dict)
|
||||
assert isinstance(load_defaults(), dict)
|
||||
|
||||
|
||||
def test_load_llm_config_temperature_zero_is_respected(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""temperature=0 是合法配置,不应被 or 链回退默认。"""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-env")
|
||||
monkeypatch.setenv("DEEPSEEK_MODEL", "deepseek-m")
|
||||
_patch_scene(monkeypatch, {"provider": "deepseek", "model": "deepseek-m", "temperature": 0})
|
||||
cfg = load_llm_config(scene="daily_report")
|
||||
assert cfg.temperature == 0.0
|
||||
|
||||
|
||||
def test_load_llm_config_scene_max_attempts_zero() -> None:
|
||||
"""max_attempts=0 由 _pick_int 显式处理。"""
|
||||
from llm.client import _pick_int
|
||||
|
||||
assert _pick_int({"max_attempts": 0}, "max_attempts", 3) == 0
|
||||
assert _pick_int({"max_attempts": ""}, "max_attempts", 3) == 3
|
||||
assert _pick_int({}, "max_attempts", 3) == 3
|
||||
@@ -0,0 +1,129 @@
|
||||
"""M8 MCP 服务测试。
|
||||
|
||||
验证工具存在 + 格式化逻辑 + 降级行为,不依赖真实嵌入/检索。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from mcp_server.tools import (
|
||||
_fmt_results,
|
||||
mcp,
|
||||
search_company_news,
|
||||
search_news,
|
||||
search_sentiment_trend,
|
||||
search_stock_events,
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 工具存在性
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_mcp_server_has_name() -> None:
|
||||
assert mcp.name == "A股DeepResearch"
|
||||
|
||||
|
||||
def test_all_five_tools_registered() -> None:
|
||||
tool_names = [getattr(t, "name", "") for t in mcp._tool_manager._tools.values()] # type: ignore[union-attr]
|
||||
expected = {
|
||||
"search_news", "search_company_news", "search_industry_news",
|
||||
"search_stock_events", "search_sentiment_trend",
|
||||
}
|
||||
assert set(tool_names) == expected
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# _fmt_results
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _hit(title: str = "测试标题", source: str = "cls", score: float = 0.9,
|
||||
sentiment: str = "positive", stock_codes: list[str] | None = None,
|
||||
company_names: list[str] | None = None, summary: str = "摘要",
|
||||
publish_time: str = "2026-06-16T10:00:00",
|
||||
url: str = "https://example.com/1") -> dict:
|
||||
return {
|
||||
"title": title, "url": url, "source": source, "score": score,
|
||||
"publish_time": publish_time,
|
||||
"event": {
|
||||
"sentiment": sentiment, "importance": 4, "event_type": "重大合同",
|
||||
"stock_codes": stock_codes or [], "company_names": company_names or [],
|
||||
"industries": ["动力电池"], "summary": summary,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_fmt_results_contains_title_and_source() -> None:
|
||||
hits = [_hit("宁德时代签百亿合同", "cls")]
|
||||
out = _fmt_results(hits, "宁德时代")
|
||||
assert "百亿合同" in out
|
||||
assert "cls" in out
|
||||
assert "0.9" in out
|
||||
|
||||
|
||||
def test_fmt_results_includes_event_fields() -> None:
|
||||
hits = [_hit(
|
||||
company_names=["宁德时代"], stock_codes=["300750"],
|
||||
)]
|
||||
out = _fmt_results(hits, "查询")
|
||||
assert "300750" in out
|
||||
assert "宁德时代" in out
|
||||
assert "动力电池" in out
|
||||
assert "摘要" in out
|
||||
|
||||
|
||||
def test_fmt_results_empty_returns_hint() -> None:
|
||||
out = _fmt_results([], "无结果")
|
||||
assert "未找到" in out and "无结果" in out
|
||||
|
||||
|
||||
def test_fmt_results_multiple_hits() -> None:
|
||||
hits = [_hit(f"测试{i}") for i in range(3)]
|
||||
out = _fmt_results(hits, "查询")
|
||||
assert "共 3 条" in out
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 工具:降级行为(无精确命中时走纯语义)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
@patch("mcp_server.tools._search")
|
||||
def test_search_company_news_falls_back_on_empty(mock_search) -> None:
|
||||
"""company 精确命中 0 条时,降级为无 filter 纯语义搜索。"""
|
||||
mock_search.side_effect = [
|
||||
[], # 第一次:精确匹配 0 条
|
||||
[_hit("fallback")], # 降级: 纯语义
|
||||
]
|
||||
out = search_company_news("查询", company="不存在的公司")
|
||||
assert "fallback" in out
|
||||
|
||||
|
||||
@patch("mcp_server.tools._search")
|
||||
def test_search_stock_events_normalizes_code(mock_search) -> None:
|
||||
"""stock_code 应去掉后缀,统一大写。"""
|
||||
mock_search.return_value = [_hit("结果")]
|
||||
out = search_stock_events("查询", stock_code="300750.SZ")
|
||||
mock_search.assert_called() # code 应为 '300750'
|
||||
assert "结果" in out
|
||||
|
||||
|
||||
@patch("mcp_server.tools._search")
|
||||
def test_search_sentiment_trend_includes_stats(mock_search) -> None:
|
||||
mock_search.return_value = [
|
||||
_hit("a", sentiment="positive"),
|
||||
_hit("b", sentiment="positive"),
|
||||
_hit("c", sentiment="negative"),
|
||||
_hit("d", sentiment="neutral"),
|
||||
]
|
||||
out = search_sentiment_trend("查询", sentiment="all", top_k=10)
|
||||
assert "利好 2" in out
|
||||
assert "利空 1" in out
|
||||
assert "中性 1" in out
|
||||
assert "共 4 条" in out
|
||||
|
||||
|
||||
@patch("mcp_server.tools._search")
|
||||
def test_search_news_passthrough(mock_search) -> None:
|
||||
mock_search.return_value = [_hit("结果")]
|
||||
out = search_news("查询")
|
||||
assert "结果" in out
|
||||
@@ -0,0 +1,285 @@
|
||||
"""日报结构化组装单元测试（_build_report_data，纯逻辑）。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import date
|
||||
|
||||
import pytest
|
||||
|
||||
from scheduler.reporter import _build_report_data
|
||||
|
||||
|
||||
def _fake_event(title: str, importance: int, event_type: str = "其他",
|
||||
sentiment: str = "neutral", source_id: str = "cls",
|
||||
url: str = "https://x.com/1") -> dict:
|
||||
return {
|
||||
"title": title,
|
||||
"url": url,
|
||||
"source_id": source_id,
|
||||
"event": {
|
||||
"stock_codes": [],
|
||||
"company_names": [],
|
||||
"industries": [],
|
||||
"sentiment": sentiment,
|
||||
"importance": importance,
|
||||
"event_type": event_type,
|
||||
"summary": f"{title}的摘要",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestBuildReportData:
|
||||
def test_sections_and_ranks(self) -> None:
|
||||
news = {
|
||||
"total": 2, "hi_threshold": 4,
|
||||
"high": [_fake_event("新闻A", 5), _fake_event("新闻B", 4)],
|
||||
"sentiments": {"neutral": 2}, "importances": {5: 1, 4: 1},
|
||||
"event_types": {"其他": 2},
|
||||
}
|
||||
cninfo = {
|
||||
"total": 1, "hi_threshold": 2,
|
||||
"high": [_fake_event("公告C", 3, event_type="公告")],
|
||||
"by_day": {"07月10日": 1}, "announcement": 1, "research": 0, "irm": 0,
|
||||
}
|
||||
pipeline = {"raw_total": 100, "proc": 90}
|
||||
xwlb = {"items": [_fake_event("联播D", 4, event_type="新闻联播", source_id="xwlb")],
|
||||
"date": "07月10日"}
|
||||
|
||||
r = _build_report_data(news, cninfo, pipeline, "AI摘要", "20260710", xwlb=xwlb)
|
||||
|
||||
assert r.report_date == date(2026, 7, 10)
|
||||
assert r.report_type == "finance"
|
||||
assert r.file_name == ""
|
||||
assert r.ai_summary == "AI摘要"
|
||||
assert [(e.section, e.rank) for e in r.events] == [
|
||||
("news", 1), ("news", 2), ("cninfo", 1), ("xwlb", 1),
|
||||
]
|
||||
assert r.events[0].source == "cls"
|
||||
assert r.events[3].source == "xwlb"
|
||||
|
||||
def test_stats_snapshot(self) -> None:
|
||||
news = {"total": 1, "hi_threshold": 4, "high": [], "sentiments": {},
|
||||
"importances": {}, "event_types": {}}
|
||||
cninfo = {"total": 0, "hi_threshold": 0, "high": [], "by_day": {},
|
||||
"announcement": 0, "research": 0, "irm": 0}
|
||||
r = _build_report_data(news, cninfo, {"raw_total": 100}, "s", "20260710")
|
||||
# stats 可 JSON 序列化（入库时 json.dumps）
|
||||
json.dumps(r.stats, ensure_ascii=False)
|
||||
assert r.stats["pipeline"] == {"raw_total": 100}
|
||||
assert r.stats["news"]["total"] == 1
|
||||
assert "xwlb" not in r.stats
|
||||
|
||||
def test_title_truncated(self) -> None:
|
||||
news = {"total": 1, "hi_threshold": 4,
|
||||
"high": [_fake_event("长" * 600, 4)], "sentiments": {},
|
||||
"importances": {}, "event_types": {}}
|
||||
cninfo = {"total": 0, "hi_threshold": 0, "high": [], "by_day": {},
|
||||
"announcement": 0, "research": 0, "irm": 0}
|
||||
r = _build_report_data(news, cninfo, {}, "s", "20260710")
|
||||
assert len(r.events[0].title) == 512
|
||||
|
||||
def test_empty_events(self) -> None:
|
||||
news = {"total": 0, "hi_threshold": 0, "high": [], "sentiments": {},
|
||||
"importances": {}, "event_types": {}}
|
||||
cninfo = {"total": 0, "hi_threshold": 0, "high": [], "by_day": {},
|
||||
"announcement": 0, "research": 0, "irm": 0}
|
||||
r = _build_report_data(news, cninfo, {}, None, "20260710")
|
||||
assert r.events == []
|
||||
assert r.ai_summary is None
|
||||
|
||||
|
||||
class TestLlmCallRetry:
|
||||
"""_llm_call 重试逻辑（纯逻辑，mock client）。"""
|
||||
|
||||
@staticmethod
|
||||
def _fake_client(failures: int):
|
||||
"""构造 mock client：前 failures 次抛 ConnectionError，之后成功。"""
|
||||
from types import SimpleNamespace
|
||||
|
||||
n = {"count": 0}
|
||||
|
||||
class Completions:
|
||||
def create(self, **kwargs):
|
||||
n["count"] += 1
|
||||
if n["count"] <= failures:
|
||||
raise ConnectionError("transient")
|
||||
return SimpleNamespace(
|
||||
choices=[SimpleNamespace(
|
||||
message=SimpleNamespace(content="今日要点摘要"),
|
||||
finish_reason="stop",
|
||||
)]
|
||||
)
|
||||
|
||||
return SimpleNamespace(chat=SimpleNamespace(completions=Completions())), n
|
||||
|
||||
@staticmethod
|
||||
def _cfg():
|
||||
from llm.client import LLMConfig
|
||||
|
||||
return LLMConfig(
|
||||
provider="deepseek", model="deepseek-v4-flash",
|
||||
api_key="sk-test", base_url="https://api.deepseek.com",
|
||||
temperature=0.3,
|
||||
)
|
||||
|
||||
def test_success_first_try(self) -> None:
|
||||
from scheduler.reporter import _llm_call
|
||||
client, n = self._fake_client(0)
|
||||
out = _llm_call(client, self._cfg(), "p")
|
||||
assert out == "今日要点摘要"
|
||||
assert n["count"] == 1
|
||||
|
||||
def test_retry_then_success(self, monkeypatch) -> None:
|
||||
import scheduler.reporter as rep
|
||||
monkeypatch.setattr(rep, "_LLM_RETRY_TIMES", 3)
|
||||
monkeypatch.setattr(rep, "_LLM_RETRY_BACKOFF_SEC", 0.01)
|
||||
client, n = self._fake_client(2) # 前 2 次失败，第 3 次成功
|
||||
out = rep._llm_call(client, self._cfg(), "p")
|
||||
assert out == "今日要点摘要"
|
||||
assert n["count"] == 3
|
||||
|
||||
def test_exhausts_retries_raises(self, monkeypatch) -> None:
|
||||
import scheduler.reporter as rep
|
||||
monkeypatch.setattr(rep, "_LLM_RETRY_TIMES", 2)
|
||||
monkeypatch.setattr(rep, "_LLM_RETRY_BACKOFF_SEC", 0.01)
|
||||
client, n = self._fake_client(99) # 一直失败
|
||||
with pytest.raises(ConnectionError):
|
||||
rep._llm_call(client, self._cfg(), "p")
|
||||
assert n["count"] == 2 # 重试 2 次后放弃
|
||||
|
||||
def test_llm_summarize_single_chunk_passes_config(self) -> None:
|
||||
"""_llm_summarize → _llm_call 全链路:必须传 LLMConfig 而非 model 字符串。
|
||||
|
||||
回归保护:生产日报曾因 _llm_call 收到 str 报
|
||||
'str' object has no attribute 'model'。
|
||||
"""
|
||||
from scheduler.reporter import _llm_summarize
|
||||
client, n = self._fake_client(0)
|
||||
out = _llm_summarize(client, self._cfg(), ["- 新闻A", "- 新闻B"], "20260812")
|
||||
assert out == "今日要点摘要"
|
||||
assert n["count"] == 1
|
||||
|
||||
def test_llm_summarize_multi_chunk_passes_config(self) -> None:
|
||||
"""多分块场景:每块 + 合并各调用一次 _llm_call,均传 config。"""
|
||||
from scheduler.reporter import _llm_summarize
|
||||
client, n = self._fake_client(0)
|
||||
# 两条长行保证触发分块
|
||||
lines = ["- " + "长新闻内容" * 300, "- " + "长新闻内容" * 300]
|
||||
out = _llm_summarize(client, self._cfg(), lines, "20260812")
|
||||
assert out == "今日要点摘要"
|
||||
assert n["count"] == 3 # 2 块 + 1 次合并
|
||||
|
||||
|
||||
class TestCollectXwlb:
|
||||
"""_collect_xwlb 取数逻辑：应查询日报前一日(已播出的联播)，并跳过内容提要。"""
|
||||
|
||||
def test_queries_previous_day_and_skips_toc(self, monkeypatch) -> None:
|
||||
import json as _json
|
||||
import urllib.request
|
||||
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
def fake_urlopen(req, timeout=15): # noqa: ARG001
|
||||
captured["url"] = req.full_url
|
||||
|
||||
class Resp:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
return False
|
||||
|
||||
def read(self):
|
||||
return _json.dumps({"data": {"news": [
|
||||
{"daily_sub_id": 1, "news_title": "内容提要", "news_days": "2026-08-04", "news_improve": "开场白"},
|
||||
{"daily_sub_id": 2, "news_title": "联播要闻A", "news_days": "2026-08-04", "news_improve": "正文A"},
|
||||
]}}).encode("utf-8")
|
||||
|
||||
return Resp()
|
||||
|
||||
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
||||
from scheduler.reporter import _collect_xwlb
|
||||
|
||||
result = _collect_xwlb("20260805")
|
||||
# 查询的是前一日(20260804)而非当日
|
||||
assert "start_date=20260804" in captured["url"]
|
||||
assert "end_date=20260804" in captured["url"]
|
||||
# 跳过第 1 条内容提要
|
||||
assert len(result["items"]) == 1
|
||||
assert result["items"][0]["title"] == "联播要闻A"
|
||||
assert result["source_date"] == "20260804"
|
||||
assert result["date"] == "08月04日"
|
||||
|
||||
|
||||
class TestCollectNewsEventsLookback:
|
||||
"""_collect_news_events 30 小时回溯逻辑。"""
|
||||
|
||||
def test_filters_30h_and_excludes_cninfo(self, monkeypatch) -> None:
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import scheduler.reporter as rep
|
||||
|
||||
now = datetime.now().astimezone()
|
||||
|
||||
def fake_load(day_str: str) -> list[dict]: # noqa: ARG001
|
||||
def ev(title: str, hours_ago: float | None, source: str = "cls",
|
||||
importance: int = 5, aware: bool = False) -> dict:
|
||||
pt = None
|
||||
if hours_ago is not None:
|
||||
t = now - timedelta(hours=hours_ago)
|
||||
pt = t.isoformat() if not aware else t.astimezone().isoformat()
|
||||
return {
|
||||
"title": title, "url": "u", "source_id": source,
|
||||
"publish_time": pt,
|
||||
"event": {"importance": importance, "sentiment": "neutral",
|
||||
"event_type": "其他", "summary": "s"},
|
||||
}
|
||||
|
||||
return [
|
||||
ev("窗口内新闻", 10),
|
||||
ev("窗口内新闻带时区", 12, aware=True),
|
||||
ev("窗口外旧闻", 40),
|
||||
ev("无时间戳", None),
|
||||
ev("公告排除", 5, source="cninfo", importance=2),
|
||||
]
|
||||
|
||||
monkeypatch.setattr(rep, "_load_events_from_dir", fake_load)
|
||||
result = rep._collect_news_events("20260805")
|
||||
|
||||
# 两个日期目录各返回 5 条(共 10): 旧闻×2、公告×2 被滤, 保留 6 条
|
||||
assert result["total"] == 6
|
||||
titles = {e["title"] for e in result["high"]}
|
||||
assert "窗口内新闻" in titles
|
||||
assert "窗口内新闻带时区" in titles
|
||||
assert "无时间戳" in titles
|
||||
assert "窗口外旧闻" not in titles
|
||||
assert "公告排除" not in titles
|
||||
|
||||
|
||||
class TestLoadEventsSources:
|
||||
"""_load_events_from_dir 读取多源(sources)字段。"""
|
||||
|
||||
def test_loads_sources_with_fallback(self, monkeypatch, tmp_path) -> None:
|
||||
import json as _json
|
||||
|
||||
import scheduler.reporter as rep
|
||||
|
||||
ev_dir = tmp_path / "data" / "events" / "20260812"
|
||||
ev_dir.mkdir(parents=True)
|
||||
(ev_dir / "aaa.json").write_text(_json.dumps({
|
||||
"source_id": "cls", "sources": ["cls", "sina", "eastmoney"],
|
||||
"title": "多源新闻", "url": "https://x/1", "event": {},
|
||||
}, ensure_ascii=False), encoding="utf-8")
|
||||
(ev_dir / "bbb.json").write_text(_json.dumps({
|
||||
"source_id": "cls", "title": "旧产物无 sources", "url": "https://x/2",
|
||||
"event": {},
|
||||
}, ensure_ascii=False), encoding="utf-8")
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
events = rep._load_events_from_dir("20260812")
|
||||
by_url = {e["url"]: e for e in events}
|
||||
# 新产物:多源完整透传
|
||||
assert by_url["https://x/1"]["sources"] == ["cls", "sina", "eastmoney"]
|
||||
# 旧产物:兜底 [source_id]
|
||||
assert by_url["https://x/2"]["sources"] == ["cls"]
|
||||
@@ -0,0 +1,58 @@
|
||||
"""日报数据模型单元测试。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date, datetime
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from report_db.models import EventRow, ReportData
|
||||
|
||||
|
||||
class TestEventRow:
|
||||
def test_minimal(self) -> None:
|
||||
ev = EventRow(section="news", rank=1, title="标题")
|
||||
assert ev.importance is None
|
||||
assert ev.sentiment is None
|
||||
|
||||
def test_full(self) -> None:
|
||||
ev = EventRow(
|
||||
section="intl", rank=2, importance=4, event_type="地缘政治",
|
||||
title="t", summary="s", sentiment="negative", source="ForexLive",
|
||||
sources=["ForexLive", "新浪财经"],
|
||||
url="https://x.com/1",
|
||||
)
|
||||
assert ev.sentiment == "negative"
|
||||
assert ev.sources == ["ForexLive", "新浪财经"]
|
||||
|
||||
def test_sources_optional(self) -> None:
|
||||
"""sources 为可选项(旧数据无多源记录)。"""
|
||||
ev = EventRow(section="news", rank=1, title="t")
|
||||
assert ev.sources is None
|
||||
|
||||
def test_missing_title_raises(self) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
EventRow(section="news", rank=1) # type: ignore[call-arg]
|
||||
|
||||
|
||||
class TestReportData:
|
||||
def test_defaults(self) -> None:
|
||||
r = ReportData(
|
||||
report_date=date(2026, 7, 11),
|
||||
report_type="finance",
|
||||
generated_at=datetime(2026, 7, 11, 7, 0),
|
||||
)
|
||||
assert r.file_name == ""
|
||||
assert r.stats == {}
|
||||
assert r.events == []
|
||||
|
||||
def test_with_events(self) -> None:
|
||||
r = ReportData(
|
||||
report_date=date(2026, 7, 11),
|
||||
report_type="finance",
|
||||
generated_at=datetime(2026, 7, 11, 7, 0),
|
||||
ai_summary="摘要",
|
||||
events=[EventRow(section="xwlb", rank=1, title="t")],
|
||||
)
|
||||
assert len(r.events) == 1
|
||||
@@ -0,0 +1,32 @@
|
||||
"""历史导入器单元测试（文件匹配/扫描逻辑，不依赖真实 DB）。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from report_import.importer import _match_file
|
||||
|
||||
|
||||
class TestMatchFile:
|
||||
def test_finance(self) -> None:
|
||||
p = Path("20260711/finance_news_daily_20260710_0720.html")
|
||||
assert _match_file(p, None, None)
|
||||
assert _match_file(p, "20260710", None)
|
||||
assert _match_file(p, None, "finance")
|
||||
assert not _match_file(p, "20260711", None) # 文件名日期不含 20260711
|
||||
assert not _match_file(p, None, "intl")
|
||||
|
||||
def test_no_timestamp_suffix(self) -> None:
|
||||
# 早期文件无时间戳后缀，也应匹配
|
||||
p = Path("20260616/finance_news_daily_20260616.html")
|
||||
assert _match_file(p, None, None)
|
||||
assert _match_file(p, "20260616", "finance")
|
||||
|
||||
def test_intl(self) -> None:
|
||||
p = Path("20260711/intl_news_daily_20260711_070304.html")
|
||||
assert _match_file(p, "20260711", "intl")
|
||||
assert not _match_file(p, None, "finance")
|
||||
|
||||
def test_non_report_ignored(self) -> None:
|
||||
assert not _match_file(Path("20260711/002714.SZ_0724.html"), None, None)
|
||||
assert not _match_file(Path("20260711/readme.md"), None, None)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""历史日报解析器单元测试（基于真实样例 HTML）。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date, datetime
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from report_import.parser import (
|
||||
ReportParseError,
|
||||
parse_finance_report,
|
||||
parse_intl_report,
|
||||
parse_report,
|
||||
)
|
||||
|
||||
FIXTURES = Path(__file__).parent / "fixtures"
|
||||
FINANCE_HTML = (FIXTURES / "finance_news_daily_20260710_0720.html").read_text(encoding="utf-8")
|
||||
INTL_HTML = (FIXTURES / "intl_news_daily_20260711_070304.html").read_text(encoding="utf-8")
|
||||
|
||||
|
||||
class TestFinanceParse:
|
||||
def test_metadata(self) -> None:
|
||||
r = parse_finance_report(FINANCE_HTML, "finance_news_daily_20260710_0720.html")
|
||||
assert r.report_date == date(2026, 7, 10)
|
||||
assert r.report_type == "finance"
|
||||
assert r.file_name == "finance_news_daily_20260710_0720.html"
|
||||
assert r.generated_at == datetime(2026, 7, 11, 7, 20, 27) # header"生成于"优先
|
||||
|
||||
def test_ai_summary_lines(self) -> None:
|
||||
r = parse_finance_report(FINANCE_HTML, "finance_news_daily_20260710_0720.html")
|
||||
assert r.ai_summary is not None
|
||||
assert len(r.ai_summary.splitlines()) >= 5
|
||||
assert "碳达峰" in r.ai_summary
|
||||
|
||||
def test_sections_and_ranks(self) -> None:
|
||||
r = parse_finance_report(FINANCE_HTML, "finance_news_daily_20260710_0720.html")
|
||||
sections = {e.section for e in r.events}
|
||||
assert sections == {"xwlb", "news", "cninfo"}
|
||||
assert sum(1 for e in r.events if e.section == "xwlb") == 16
|
||||
assert sum(1 for e in r.events if e.section == "news") == 20
|
||||
assert sum(1 for e in r.events if e.section == "cninfo") == 20
|
||||
|
||||
def test_event_fields(self) -> None:
|
||||
r = parse_finance_report(FINANCE_HTML, "finance_news_daily_20260710_0720.html")
|
||||
xwlb = next(e for e in r.events if e.section == "xwlb")
|
||||
assert xwlb.importance == 4
|
||||
assert xwlb.event_type == "新闻联播"
|
||||
assert xwlb.sentiment == "neutral"
|
||||
assert "张国清" in xwlb.title
|
||||
news = next(e for e in r.events if e.section == "news" and e.source)
|
||||
assert news.source # 新闻板块带来源
|
||||
|
||||
def test_stats_keys(self) -> None:
|
||||
r = parse_finance_report(FINANCE_HTML, "finance_news_daily_20260710_0720.html")
|
||||
for key in ("pipeline", "sources", "sentiment", "importance", "event_types"):
|
||||
assert key in r.stats, f"缺少 stats.{key}"
|
||||
assert r.stats["pipeline"]["M1 原始文章"] == 758
|
||||
|
||||
|
||||
class TestIntlParse:
|
||||
def test_metadata(self) -> None:
|
||||
r = parse_intl_report(INTL_HTML, "intl_news_daily_20260711_070304.html")
|
||||
assert r.report_date == date(2026, 7, 11)
|
||||
assert r.report_type == "intl"
|
||||
assert r.generated_at == datetime(2026, 7, 11, 7, 3, 30)
|
||||
|
||||
def test_events_and_source_from_small(self) -> None:
|
||||
r = parse_intl_report(INTL_HTML, "intl_news_daily_20260711_070304.html")
|
||||
assert len(r.events) == 19
|
||||
first = r.events[0]
|
||||
assert first.section == "intl"
|
||||
assert first.source == "investinglive.com" # 从摘要 <small>[来源]</small> 提取
|
||||
assert first.url.startswith("https://investinglive.com/")
|
||||
assert first.importance == 4
|
||||
assert first.event_type == "地缘政治"
|
||||
# 摘要中不应残留 [来源] 标记
|
||||
assert first.summary is not None and "[investinglive.com]" not in first.summary
|
||||
|
||||
def test_stats_keys(self) -> None:
|
||||
r = parse_intl_report(INTL_HTML, "intl_news_daily_20260711_070304.html")
|
||||
for key in ("pipeline", "sentiment", "importance", "event_types", "source_dist"):
|
||||
assert key in r.stats
|
||||
assert r.stats["source_dist"][0]["来源"] == "ForexLive"
|
||||
|
||||
|
||||
class TestParseReportDispatch:
|
||||
def test_dispatch_finance(self) -> None:
|
||||
r = parse_report(FINANCE_HTML, "finance_news_daily_20260710_0720.html")
|
||||
assert r.report_type == "finance"
|
||||
|
||||
def test_dispatch_intl(self) -> None:
|
||||
r = parse_report(INTL_HTML, "intl_news_daily_20260711_070304.html")
|
||||
assert r.report_type == "intl"
|
||||
|
||||
def test_bad_filename_raises(self) -> None:
|
||||
with pytest.raises(ReportParseError):
|
||||
parse_report("<html></html>", "not_a_report.html")
|
||||
@@ -0,0 +1,132 @@
|
||||
"""M7 定时任务测试。
|
||||
|
||||
使用 mock subprocess.run,不依赖真实脚本执行。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from scheduler import PipelineResult, StepResult, run_pipeline, run_step
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# fixtures
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _mock_proc(returncode: int = 0, stderr: str = "") -> MagicMock:
|
||||
p = MagicMock()
|
||||
p.returncode = returncode
|
||||
p.stderr = stderr
|
||||
p.stdout = "ok"
|
||||
return p
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# StepResult 模型
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_step_result_defaults() -> None:
|
||||
sr = StepResult(name="crawler", success=True, elapsed_sec=12.3)
|
||||
assert sr.name == "crawler"
|
||||
assert sr.success is True
|
||||
assert sr.elapsed_sec == 12.3
|
||||
assert sr.exit_code is None
|
||||
|
||||
|
||||
def test_pipeline_result_all_success() -> None:
|
||||
pr = PipelineResult(steps=[
|
||||
StepResult(name="crawler", success=True, elapsed_sec=1),
|
||||
StepResult(name="extractor", success=True, elapsed_sec=2),
|
||||
])
|
||||
assert pr.all_success is True
|
||||
|
||||
|
||||
def test_pipeline_result_partial_failure() -> None:
|
||||
pr = PipelineResult(steps=[
|
||||
StepResult(name="crawler", success=True, elapsed_sec=1),
|
||||
StepResult(name="extractor", success=False, elapsed_sec=0),
|
||||
])
|
||||
assert pr.all_success is False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# run_step
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_run_step_success() -> None:
|
||||
with patch("subprocess.run", return_value=_mock_proc(returncode=0,
|
||||
stderr="INFO | 完成: 成功 20/20")):
|
||||
sr = run_step("crawler", "20260616")
|
||||
assert sr.success is True
|
||||
assert sr.exit_code == 0
|
||||
|
||||
|
||||
def test_run_step_failure() -> None:
|
||||
with patch("subprocess.run", return_value=_mock_proc(returncode=1)):
|
||||
sr = run_step("extractor", "20260616")
|
||||
assert sr.success is False
|
||||
assert sr.exit_code == 1
|
||||
assert "rc=1" in sr.tail_msg
|
||||
|
||||
|
||||
def test_run_step_timeout() -> None:
|
||||
with patch("subprocess.run", side_effect=subprocess.TimeoutExpired(cmd=["uv"], timeout=10)):
|
||||
sr = run_step("llm", "20260616")
|
||||
assert sr.success is False
|
||||
assert "超时" in sr.tail_msg
|
||||
|
||||
|
||||
def test_run_step_exception() -> None:
|
||||
with patch("subprocess.run", side_effect=OSError("磁盘满")):
|
||||
sr = run_step("embedding", "20260616")
|
||||
assert sr.success is False
|
||||
assert "磁盘满" in sr.tail_msg
|
||||
|
||||
|
||||
def test_run_step_unknown_name() -> None:
|
||||
sr = run_step("nonexistent", "20260616")
|
||||
assert sr.success is False
|
||||
assert "未知" in sr.tail_msg
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# run_pipeline
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_run_pipeline_all_success() -> None:
|
||||
with (
|
||||
patch("subprocess.run", return_value=_mock_proc(returncode=0)),
|
||||
patch("scheduler.reporter.generate_report", return_value=Path("/tmp/r.html")),
|
||||
):
|
||||
result = run_pipeline("20260616", steps=["crawler", "extractor", "dedup"])
|
||||
assert len(result.steps) == 3
|
||||
assert result.all_success is True
|
||||
assert result.started_at is not None
|
||||
assert result.finished_at is not None
|
||||
|
||||
|
||||
def test_run_pipeline_continues_on_failure() -> None:
|
||||
"""中间步骤失败,后续继续执行(不阻断)。"""
|
||||
call_count = {"n": 0}
|
||||
|
||||
def _side_effect(*args, **kwargs):
|
||||
call_count["n"] += 1
|
||||
if call_count["n"] == 2: # extractor 失败
|
||||
return _mock_proc(returncode=1, stderr="GNE 提取异常")
|
||||
return _mock_proc(returncode=0, stderr="ok")
|
||||
|
||||
with patch("subprocess.run", side_effect=_side_effect):
|
||||
result = run_pipeline("20260616", steps=["crawler", "extractor", "dedup", "llm"])
|
||||
assert len(result.steps) == 4
|
||||
# extractor 失败,但后续仍执行
|
||||
assert result.steps[1].success is False
|
||||
assert result.steps[2].success is True
|
||||
|
||||
|
||||
def test_run_pipeline_custom_steps() -> None:
|
||||
with patch("subprocess.run", return_value=_mock_proc(returncode=0, stderr="ok")):
|
||||
result = run_pipeline("20260616", steps=["crawler", "extractor"])
|
||||
assert len(result.steps) == 2
|
||||
assert result.all_success is True
|
||||
@@ -0,0 +1,243 @@
|
||||
"""M6 Qdrant 向量存储模块测试。
|
||||
|
||||
使用 qdrant-client 内存模式(:memory:),不依赖 Docker。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vectorstore import (
|
||||
SearchFilter,
|
||||
SearchResult,
|
||||
VectorStore,
|
||||
make_qdrant_client,
|
||||
)
|
||||
from vectorstore.client import _build_filter
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# fixtures
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
@pytest.fixture
|
||||
def store() -> VectorStore:
|
||||
c = make_qdrant_client(memory=True)
|
||||
s = VectorStore(c, collection_name="test_m6", vector_dim=4)
|
||||
s.init_collection()
|
||||
yield s
|
||||
s.close()
|
||||
|
||||
|
||||
def _point(id_: str, vector: list[float], **payload: object) -> dict:
|
||||
return {"id": id_, "vector": vector, "payload": dict(payload)}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Collection 管理
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_init_collection_creates(store: VectorStore) -> None:
|
||||
info = store.info()
|
||||
assert info.exists is True
|
||||
assert info.name == "test_m6"
|
||||
|
||||
|
||||
def test_init_collection_idempotent(store: VectorStore) -> None:
|
||||
"""再次 init 不应报错,count 不变。"""
|
||||
store.upsert([_point("a", [1.0, 0, 0, 0])])
|
||||
store.init_collection() # 不应重建
|
||||
assert store.count() == 1
|
||||
|
||||
|
||||
def test_init_collection_recreate_clears(store: VectorStore) -> None:
|
||||
store.upsert([_point("a", [1.0, 0, 0, 0])])
|
||||
store.init_collection(recreate=True)
|
||||
assert store.count() == 0
|
||||
|
||||
|
||||
def test_delete_collection(store: VectorStore) -> None:
|
||||
store.delete_collection()
|
||||
assert store.info().exists is False
|
||||
# 再次 init 应恢复
|
||||
store.init_collection()
|
||||
assert store.info().exists is True
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# upsert + count
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_upsert_and_count(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a", [1, 0, 0, 0], title="Article A"),
|
||||
_point("b", [0, 1, 0, 0], title="Article B"),
|
||||
])
|
||||
assert store.count() == 2
|
||||
|
||||
|
||||
def test_upsert_idempotent(store: VectorStore) -> None:
|
||||
"""同 url_hash 再次 upsert 不应增加 count,数据被覆盖。"""
|
||||
store.upsert([_point("a", [1, 0, 0, 0], title="Old")])
|
||||
store.upsert([_point("a", [0, 0, 0, 1], title="New")])
|
||||
assert store.count() == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# query - 语义检索
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_query_returns_score_desc(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a", [1.0, 0, 0, 0], title="A"),
|
||||
_point("b", [0.0, 1.0, 0, 0], title="B"),
|
||||
_point("c", [0.0, 0, 1.0, 0], title="C"),
|
||||
])
|
||||
results = store.query(query_vector=[0.9, 0.1, 0, 0], top_k=2)
|
||||
assert len(results) == 2
|
||||
assert results[0].url_hash == "a"
|
||||
# score 应递减
|
||||
assert results[0].score >= results[1].score
|
||||
|
||||
|
||||
def test_query_score_threshold(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a", [1, 0, 0, 0], title="A"),
|
||||
_point("b", [0, 1, 0, 0], title="B"),
|
||||
])
|
||||
# 只有 a 会匹配
|
||||
results = store.query(query_vector=[1, 0, 0, 0], top_k=10, score_threshold=0.9)
|
||||
assert len(results) == 1
|
||||
assert results[0].url_hash == "a"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# query - 结构化过滤
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_query_filter_by_source_id(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a1", [1, 0, 0, 0], source_id="cls", title="CLS article"),
|
||||
_point("a2", [0.9, 0.1, 0, 0], source_id="sina", title="Sina article"),
|
||||
])
|
||||
results = store.query(
|
||||
query_vector=[1, 0, 0, 0],
|
||||
filter=SearchFilter(source_id="cls"),
|
||||
top_k=5,
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].source_id == "cls"
|
||||
|
||||
|
||||
def test_query_filter_by_stock_codes(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a", [1, 0, 0, 0], source_id="cls",
|
||||
event={"stock_codes": ["300750"], "sentiment": "positive"}),
|
||||
_point("b", [0.9, 0.1, 0, 0], source_id="sina",
|
||||
event={"stock_codes": ["000001"], "sentiment": "neutral"}),
|
||||
_point("c", [0.8, 0.2, 0, 0], source_id="sina",
|
||||
event={"stock_codes": ["300750"], "sentiment": "negative"}),
|
||||
])
|
||||
results = store.query(
|
||||
query_vector=[1, 0, 0, 0],
|
||||
filter=SearchFilter(stock_codes=["300750"]),
|
||||
top_k=5,
|
||||
)
|
||||
assert len(results) == 2
|
||||
for r in results:
|
||||
assert "300750" in (r.event or {}).get("stock_codes", [])
|
||||
|
||||
|
||||
def test_query_filter_by_sentiment(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a", [1, 0, 0, 0], source_id="cls",
|
||||
event={"sentiment": "positive"}),
|
||||
_point("b", [0, 1, 0, 0], source_id="cls",
|
||||
event={"sentiment": "negative"}),
|
||||
])
|
||||
results = store.query(
|
||||
query_vector=[1, 0, 0, 0],
|
||||
filter=SearchFilter(sentiment="positive"),
|
||||
top_k=5,
|
||||
)
|
||||
assert len(results) >= 1
|
||||
assert all((r.event or {}).get("sentiment") == "positive" for r in results)
|
||||
|
||||
|
||||
def test_query_filter_by_importance_min(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a", [1, 0, 0, 0], source_id="cls",
|
||||
event={"importance": 2}),
|
||||
_point("b", [0, 1, 0, 0], source_id="cls",
|
||||
event={"importance": 4}),
|
||||
_point("c", [0, 0, 1, 0], source_id="cls",
|
||||
event={"importance": 5}),
|
||||
])
|
||||
results = store.query(
|
||||
query_vector=[0.5, 0.5, 0.5, 0],
|
||||
filter=SearchFilter(importance_min=4),
|
||||
top_k=5,
|
||||
)
|
||||
assert all((r.event or {}).get("importance", 0) >= 4 for r in results)
|
||||
|
||||
|
||||
def test_query_filter_by_industry(store: VectorStore) -> None:
|
||||
store.upsert([
|
||||
_point("a", [1, 0, 0, 0], source_id="cls",
|
||||
event={"industries": ["动力电池"]}),
|
||||
_point("b", [0, 1, 0, 0], source_id="cls",
|
||||
event={"industries": ["白酒"]}),
|
||||
])
|
||||
results = store.query(
|
||||
query_vector=[1, 0, 0, 0],
|
||||
filter=SearchFilter(industries=["动力电池"]),
|
||||
top_k=5,
|
||||
)
|
||||
assert len(results) >= 1
|
||||
for r in results:
|
||||
assert "动力电池" in (r.event or {}).get("industries", [])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Filter 构建
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_build_filter_empty_returns_none() -> None:
|
||||
assert _build_filter(SearchFilter()) is None
|
||||
|
||||
|
||||
def test_build_filter_source_id() -> None:
|
||||
f = _build_filter(SearchFilter(source_id="cls"))
|
||||
assert f is not None and len(f.must) == 1 # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_build_filter_date_range() -> None:
|
||||
f = _build_filter(SearchFilter(publish_date_from="2026-06-01", publish_date_to="2026-06-30"))
|
||||
assert f is not None
|
||||
# range 应含 gte + lte
|
||||
cond = f.must[0] # type: ignore[union-attr]
|
||||
assert cond.key == "publish_time"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# search_result 模型
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_search_result_short_summary() -> None:
|
||||
r = SearchResult(
|
||||
url_hash="abc",
|
||||
score=0.95,
|
||||
title="宁德时代签约 100GWh 协议",
|
||||
url="https://x/1",
|
||||
source_id="cls",
|
||||
event={"stock_codes": ["300750"], "sentiment": "positive"},
|
||||
)
|
||||
s = r.short_summary()
|
||||
assert "cls" in s and "0.9500" in s and "300750" in s
|
||||
|
||||
|
||||
def test_search_result_handles_none_event() -> None:
|
||||
r = SearchResult(
|
||||
url_hash="abc", score=0.5, title="t", url="u", source_id="cls", event=None,
|
||||
)
|
||||
s = r.short_summary()
|
||||
assert "-" in s # 无 stock_code
|
||||
Reference in New Issue
Block a user