Initial commit
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
|
||||
@@ -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,382 @@
|
||||
"""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
|
||||
@@ -0,0 +1,376 @@
|
||||
"""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_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,351 @@
|
||||
"""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
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 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.delenv("LLM_MODEL", raising=False)
|
||||
cfg = load_llm_config()
|
||||
assert cfg.provider == "deepseek"
|
||||
assert cfg.api_key == "sk-test-deepseek"
|
||||
assert cfg.model.startswith("deepseek")
|
||||
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.delenv("LLM_MODEL", raising=False)
|
||||
cfg = load_llm_config()
|
||||
assert cfg.provider == "qwen"
|
||||
assert cfg.api_key == "sk-test-qwen"
|
||||
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)
|
||||
with pytest.raises(ValueError):
|
||||
load_llm_config(provider="deepseek")
|
||||
@@ -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,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