初始化
This commit is contained in:
@@ -0,0 +1,622 @@
|
||||
"""M4 LLM 翻译 + 事件抽取模块单元测试。"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from openai import OpenAI
|
||||
|
||||
from extractor.models import ProcessedArticle
|
||||
from llm.client import LLMConfig, load_llm_config
|
||||
from llm.extractor import (
|
||||
MAX_CONTENT_CHARS,
|
||||
PromptTemplate,
|
||||
_extract_json_object,
|
||||
parse_translation_json,
|
||||
translate_and_extract,
|
||||
)
|
||||
from llm.models import (
|
||||
INTERNATIONAL_EVENT_TYPES,
|
||||
EnTranslatedArticle,
|
||||
EventExtraction,
|
||||
LLMCallError,
|
||||
LLMTranslationOutput,
|
||||
Sentiment,
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 辅助工厂
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _make_article(
|
||||
*,
|
||||
url: str = "https://www.reuters.com/business/1",
|
||||
url_hash: str = "abc1234567890000",
|
||||
source_id: str = "reuters",
|
||||
source_name: str = "Reuters",
|
||||
title: str = "Fed Holds Rates Steady as Markets Rally",
|
||||
content: str = (
|
||||
"The Federal Reserve held interest rates steady on Wednesday, "
|
||||
"citing solid economic growth and a strong labor market. "
|
||||
"Markets rallied in response, with the S&P 500 gaining 1.2 percent."
|
||||
),
|
||||
publish_time: str = "2026-06-16T10:00:00",
|
||||
word_count: int = 0,
|
||||
) -> ProcessedArticle:
|
||||
return ProcessedArticle(
|
||||
source_id=source_id,
|
||||
source_name=source_name,
|
||||
url=url,
|
||||
url_hash=url_hash,
|
||||
title=title,
|
||||
content=content,
|
||||
publish_time=publish_time,
|
||||
word_count=word_count or len(content.split()),
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Sentiment / 枚举
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestSentiment:
|
||||
"""Sentiment 枚举测试。"""
|
||||
|
||||
def test_values(self):
|
||||
assert Sentiment.POSITIVE == "positive"
|
||||
assert Sentiment.NEUTRAL == "neutral"
|
||||
assert Sentiment.NEGATIVE == "negative"
|
||||
|
||||
def test_from_string(self):
|
||||
assert Sentiment("positive") == Sentiment.POSITIVE
|
||||
assert Sentiment("neutral") == Sentiment.NEUTRAL
|
||||
assert Sentiment("negative") == Sentiment.NEGATIVE
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# EventExtraction 模型校验
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestEventExtraction:
|
||||
"""EventExtraction 模型测试。"""
|
||||
|
||||
def test_valid_event(self):
|
||||
ev = EventExtraction(
|
||||
event_type="财报披露",
|
||||
stock_codes=["AAPL", "TSLA"],
|
||||
sentiment="positive",
|
||||
importance=4,
|
||||
summary_zh="苹果第三季度营收超预期",
|
||||
)
|
||||
assert ev.event_type == "财报披露"
|
||||
assert ev.stock_codes == ["AAPL", "TSLA"]
|
||||
assert ev.sentiment == Sentiment.POSITIVE
|
||||
assert ev.importance == 4
|
||||
|
||||
def test_stock_codes_filtered_and_uppercased(self):
|
||||
ev = EventExtraction(
|
||||
event_type="并购收购",
|
||||
stock_codes=["aapl", " msft ", "", "INVALID123", "GOOGL"],
|
||||
sentiment="neutral",
|
||||
importance=3,
|
||||
summary_zh="微软收购测试",
|
||||
)
|
||||
# INVALID123 > 5 chars → 过滤；空 → 过滤；小写 → 大写；去重
|
||||
assert ev.stock_codes == ["AAPL", "MSFT", "GOOGL"]
|
||||
|
||||
def test_empty_stock_codes(self):
|
||||
ev = EventExtraction(
|
||||
event_type="宏观经济",
|
||||
stock_codes=[],
|
||||
sentiment="neutral",
|
||||
importance=2,
|
||||
summary_zh="GDP 数据发布",
|
||||
)
|
||||
assert ev.stock_codes == []
|
||||
|
||||
def test_importance_bounds(self):
|
||||
# 1 和 5 都能通过
|
||||
ev1 = EventExtraction(
|
||||
event_type="宏观经济", sentiment="neutral", importance=1, summary_zh="t1"
|
||||
)
|
||||
assert ev1.importance == 1
|
||||
ev5 = EventExtraction(
|
||||
event_type="央行决议", sentiment="negative", importance=5, summary_zh="t5"
|
||||
)
|
||||
assert ev5.importance == 5
|
||||
|
||||
def test_importance_out_of_range_rejected(self):
|
||||
with pytest.raises(Exception):
|
||||
EventExtraction(
|
||||
event_type="其他", sentiment="neutral", importance=0, summary_zh="t"
|
||||
)
|
||||
with pytest.raises(Exception):
|
||||
EventExtraction(
|
||||
event_type="其他", sentiment="neutral", importance=6, summary_zh="t"
|
||||
)
|
||||
|
||||
def test_invalid_sentiment_rejected(self):
|
||||
with pytest.raises(Exception):
|
||||
EventExtraction(
|
||||
event_type="其他", stock_codes=[], sentiment="happy", importance=3,
|
||||
summary_zh="t",
|
||||
)
|
||||
|
||||
def test_event_type_normalized(self):
|
||||
"""空 event_type 默认"其他"。"""
|
||||
ev = EventExtraction(
|
||||
event_type="", sentiment="neutral", importance=2, summary_zh="测试"
|
||||
)
|
||||
assert ev.event_type == "其他"
|
||||
|
||||
def test_summary_zh_max_length_rejected(self):
|
||||
"""超长 summary_zh（>200 字符）直接拒绝。"""
|
||||
long_summary = "测试" * 150 # 300 chars > 200
|
||||
with pytest.raises(Exception):
|
||||
EventExtraction(
|
||||
event_type="行业动态",
|
||||
sentiment="neutral",
|
||||
importance=2,
|
||||
summary_zh=long_summary,
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# LLMTranslationOutput
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestLLMTranslationOutput:
|
||||
"""LLMTranslationOutput 模型测试。"""
|
||||
|
||||
def test_valid_full_output(self):
|
||||
data = {
|
||||
"title_zh": "美联储维持利率不变，市场上涨",
|
||||
"content_zh": "美联储周三维持利率不变...",
|
||||
"events": [
|
||||
{
|
||||
"event_type": "央行决议",
|
||||
"stock_codes": [],
|
||||
"sentiment": "positive",
|
||||
"importance": 5,
|
||||
"summary_zh": "美联储维持利率不变",
|
||||
}
|
||||
],
|
||||
}
|
||||
out = LLMTranslationOutput.model_validate(data)
|
||||
assert out.title_zh == data["title_zh"]
|
||||
assert len(out.events) == 1
|
||||
assert out.events[0].event_type == "央行决议"
|
||||
|
||||
def test_no_events(self):
|
||||
data = {
|
||||
"title_zh": "每日市场简报",
|
||||
"content_zh": "今日市场整体平淡...",
|
||||
"events": [],
|
||||
}
|
||||
out = LLMTranslationOutput.model_validate(data)
|
||||
assert out.events == []
|
||||
|
||||
def test_missing_title_zh_rejected(self):
|
||||
data = {
|
||||
"content_zh": "正文...",
|
||||
"events": [],
|
||||
}
|
||||
with pytest.raises(Exception):
|
||||
LLMTranslationOutput.model_validate(data)
|
||||
|
||||
def test_missing_content_zh_rejected(self):
|
||||
data = {
|
||||
"title_zh": "标题",
|
||||
"events": [],
|
||||
}
|
||||
with pytest.raises(Exception):
|
||||
LLMTranslationOutput.model_validate(data)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# EnTranslatedArticle
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestEnTranslatedArticle:
|
||||
"""EnTranslatedArticle 模型测试。"""
|
||||
|
||||
def test_minimal_construction(self):
|
||||
article = EnTranslatedArticle(
|
||||
source_id="reuters",
|
||||
source_name="Reuters",
|
||||
url="https://example.com/1",
|
||||
url_hash="abc123",
|
||||
title="Fed Holds Rates",
|
||||
title_zh="美联储维持利率",
|
||||
content_en="The Fed held rates steady.",
|
||||
content_zh="美联储维持利率不变。",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
)
|
||||
assert article.events == []
|
||||
assert article.word_count_zh == 0
|
||||
|
||||
def test_with_events(self):
|
||||
article = EnTranslatedArticle(
|
||||
source_id="reuters",
|
||||
source_name="Reuters",
|
||||
url="https://example.com/1",
|
||||
url_hash="abc123",
|
||||
title="Apple Earnings",
|
||||
title_zh="苹果财报",
|
||||
content_en="Apple reported record earnings.",
|
||||
content_zh="苹果公布了创纪录的财报。",
|
||||
events=[
|
||||
EventExtraction(
|
||||
event_type="财报披露",
|
||||
stock_codes=["AAPL"],
|
||||
sentiment="positive",
|
||||
importance=4,
|
||||
summary_zh="苹果财报超预期",
|
||||
)
|
||||
],
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
)
|
||||
assert len(article.events) == 1
|
||||
assert "AAPL" in article.short_summary()
|
||||
|
||||
def test_short_summary_no_events(self):
|
||||
article = EnTranslatedArticle(
|
||||
source_id="reuters",
|
||||
source_name="Reuters",
|
||||
url="https://example.com/1",
|
||||
url_hash="abc123",
|
||||
title="Market Wrap",
|
||||
title_zh="市场综述",
|
||||
content_en="Markets were flat today.",
|
||||
content_zh="今日市场持平。",
|
||||
provider="deepseek",
|
||||
model="deepseek-v4-flash",
|
||||
)
|
||||
assert "-" in article.short_summary() or "0events" in article.short_summary()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# PromptTemplate
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestPromptTemplate:
|
||||
"""PromptTemplate 测试。"""
|
||||
|
||||
def test_parse_and_render(self, tmp_path: Path):
|
||||
"""测试模板解析和渲染。"""
|
||||
prompt_content = """# 测试标题
|
||||
|
||||
## System Prompt
|
||||
|
||||
你是翻译助手。
|
||||
|
||||
---
|
||||
|
||||
## User Input
|
||||
|
||||
标题: {title}
|
||||
来源: {source_name}
|
||||
正文: {content}
|
||||
"""
|
||||
prompt_path = tmp_path / "test_prompt.md"
|
||||
prompt_path.write_text(prompt_content, encoding="utf-8")
|
||||
|
||||
tpl = PromptTemplate(template_path=prompt_path)
|
||||
article = _make_article()
|
||||
system, user = tpl.render(article)
|
||||
|
||||
assert "翻译助手" in system
|
||||
assert article.title in user
|
||||
assert article.source_name in user
|
||||
assert article.content in user
|
||||
|
||||
def test_content_truncation(self, tmp_path: Path):
|
||||
"""测试超长正文截断。"""
|
||||
prompt_content = """## System Prompt
|
||||
你是助手。
|
||||
---
|
||||
|
||||
## User Input
|
||||
正文: {content}
|
||||
"""
|
||||
prompt_path = tmp_path / "test_prompt.md"
|
||||
prompt_path.write_text(prompt_content, encoding="utf-8")
|
||||
|
||||
tpl = PromptTemplate(template_path=prompt_path)
|
||||
long_content = "X" * (MAX_CONTENT_CHARS + 500)
|
||||
article = _make_article(content=long_content)
|
||||
|
||||
_system, user = tpl.render(article)
|
||||
assert "正文过长已截断" in user
|
||||
assert len("X" * MAX_CONTENT_CHARS) + len("\n\n[正文过长已截断]") < len(user)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# JSON 提取
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestExtractJsonObject:
|
||||
"""_extract_json_object 函数测试。"""
|
||||
|
||||
def test_plain_json(self):
|
||||
raw = '{"key": "value"}'
|
||||
assert _extract_json_object(raw) == '{"key": "value"}'
|
||||
|
||||
def test_json_with_fence(self):
|
||||
raw = '```json\n{"key": "value"}\n```'
|
||||
assert _extract_json_object(raw) == '{"key": "value"}'
|
||||
|
||||
def test_json_with_text_before(self):
|
||||
raw = 'Here is the result:\n{"key": "value"}'
|
||||
assert _extract_json_object(raw) == '{"key": "value"}'
|
||||
|
||||
def test_nested_braces(self):
|
||||
raw = '{"outer": {"inner": [1, 2, 3]}}'
|
||||
assert _extract_json_object(raw) == raw.strip()
|
||||
|
||||
def test_empty_string(self):
|
||||
assert _extract_json_object("") == ""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# parse_translation_json
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestParseTranslationJson:
|
||||
"""parse_translation_json 函数测试。"""
|
||||
|
||||
def test_valid_json(self):
|
||||
raw = json.dumps({
|
||||
"title_zh": "测试标题",
|
||||
"content_zh": "测试正文",
|
||||
"events": [],
|
||||
})
|
||||
result = parse_translation_json(raw)
|
||||
assert result.title_zh == "测试标题"
|
||||
assert result.content_zh == "测试正文"
|
||||
|
||||
def test_invalid_json(self):
|
||||
with pytest.raises(LLMCallError, match="JSON 解析失败"):
|
||||
parse_translation_json("not valid json {{{")
|
||||
|
||||
def test_non_object(self):
|
||||
with pytest.raises(LLMCallError, match="非对象"):
|
||||
parse_translation_json("[1, 2, 3]")
|
||||
|
||||
def test_schema_validation_fails(self):
|
||||
"""缺少必填字段时抛出 LLMCallError。"""
|
||||
raw = json.dumps({"title_zh": "标题"}) # 缺少 content_zh
|
||||
with pytest.raises(LLMCallError, match="schema 校验失败"):
|
||||
parse_translation_json(raw)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# LLMConfig / load_llm_config
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestLLMConfig:
|
||||
"""LLMConfig 测试。"""
|
||||
|
||||
def test_valid_config(self):
|
||||
cfg = LLMConfig(
|
||||
provider="deepseek",
|
||||
model="deepseek-chat",
|
||||
api_key="sk-test",
|
||||
base_url="https://api.deepseek.com",
|
||||
)
|
||||
assert cfg.provider == "deepseek"
|
||||
|
||||
def test_empty_api_key_raises(self):
|
||||
with pytest.raises(ValueError, match="API key 为空"):
|
||||
LLMConfig(
|
||||
provider="deepseek",
|
||||
model="deepseek-chat",
|
||||
api_key="",
|
||||
base_url="https://api.deepseek.com",
|
||||
)
|
||||
|
||||
|
||||
class TestLoadLLMConfig:
|
||||
"""load_llm_config 函数测试。"""
|
||||
|
||||
def test_deepseek_from_env(self, monkeypatch):
|
||||
"""从环境变量构造 DeepSeek 配置。"""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek-test-key")
|
||||
config = load_llm_config(provider="deepseek")
|
||||
assert config.provider == "deepseek"
|
||||
assert config.api_key == "sk-deepseek-test-key"
|
||||
assert "deepseek" in config.base_url
|
||||
|
||||
def test_qwen_fallback_to_dashscope_key(self, monkeypatch):
|
||||
"""Qwen 的 API Key 可回退到 DASHSCOPE_API_KEY。"""
|
||||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-dashscope-key")
|
||||
monkeypatch.delenv("QWEN_API_KEY", raising=False)
|
||||
config = load_llm_config(provider="qwen")
|
||||
assert config.provider == "qwen"
|
||||
assert config.api_key == "sk-dashscope-key"
|
||||
|
||||
def test_unknown_provider_raises(self):
|
||||
with pytest.raises(ValueError, match="未知 LLM provider"):
|
||||
load_llm_config(provider="openai")
|
||||
|
||||
def test_missing_api_key_raises(self, monkeypatch):
|
||||
"""未配置 API Key 时应抛出明确错误。"""
|
||||
monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False)
|
||||
with pytest.raises(ValueError, match="API key 未配置"):
|
||||
load_llm_config(provider="deepseek")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# translate_and_extract（mock LLM）
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestTranslateAndExtract:
|
||||
"""translate_and_extract 测试（mock LLM 响应）。"""
|
||||
|
||||
def test_successful_translation(self, monkeypatch):
|
||||
"""Mock LLM 返回有效 JSON。"""
|
||||
config = LLMConfig(
|
||||
provider="deepseek",
|
||||
model="test-model",
|
||||
api_key="sk-test",
|
||||
base_url="https://test.api",
|
||||
)
|
||||
article = _make_article()
|
||||
|
||||
# Mock OpenAI client
|
||||
mock_client = MagicMock(spec=OpenAI)
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = json.dumps({
|
||||
"title_zh": "美联储维持利率不变",
|
||||
"content_zh": "美联储周三维持利率不变，理由是经济增长稳健。",
|
||||
"events": [
|
||||
{
|
||||
"event_type": "央行决议",
|
||||
"stock_codes": [],
|
||||
"sentiment": "positive",
|
||||
"importance": 5,
|
||||
"summary_zh": "美联储维持利率不变，市场反弹",
|
||||
}
|
||||
],
|
||||
})
|
||||
mock_response.usage = MagicMock()
|
||||
mock_response.usage.prompt_tokens = 500
|
||||
mock_response.usage.completion_tokens = 200
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
|
||||
result = translate_and_extract(
|
||||
client=mock_client,
|
||||
config=config,
|
||||
article=article,
|
||||
)
|
||||
|
||||
assert result.title_zh == "美联储维持利率不变"
|
||||
assert len(result.events) == 1
|
||||
assert result.events[0].event_type == "央行决议"
|
||||
assert result.provider == "deepseek"
|
||||
assert result.prompt_tokens == 500
|
||||
assert result.completion_tokens == 200
|
||||
assert result.word_count_zh > 0
|
||||
|
||||
def test_empty_content_zh_retries(self, monkeypatch):
|
||||
"""LLM 返回空 content_zh 时触发重试。"""
|
||||
config = LLMConfig(
|
||||
provider="deepseek",
|
||||
model="test-model",
|
||||
api_key="sk-test",
|
||||
base_url="https://test.api",
|
||||
)
|
||||
article = _make_article()
|
||||
|
||||
mock_client = MagicMock(spec=OpenAI)
|
||||
# 始终返回空 content_zh
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = json.dumps({
|
||||
"title_zh": "标题",
|
||||
"content_zh": "",
|
||||
"events": [],
|
||||
})
|
||||
mock_response.usage = MagicMock()
|
||||
mock_response.usage.prompt_tokens = 100
|
||||
mock_response.usage.completion_tokens = 10
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
|
||||
with pytest.raises(LLMCallError, match="放弃"):
|
||||
translate_and_extract(
|
||||
client=mock_client,
|
||||
config=config,
|
||||
article=article,
|
||||
max_attempts=2, # 减少重试加速测试
|
||||
)
|
||||
|
||||
def test_retry_on_json_parse_failure(self, monkeypatch):
|
||||
"""前 N-1 次返回无效 JSON，最后一次成功。"""
|
||||
config = LLMConfig(
|
||||
provider="deepseek",
|
||||
model="test-model",
|
||||
api_key="sk-test",
|
||||
base_url="https://test.api",
|
||||
)
|
||||
article = _make_article()
|
||||
|
||||
mock_client = MagicMock(spec=OpenAI)
|
||||
# 第一次失败，第二次成功
|
||||
mock_client.chat.completions.create.side_effect = [
|
||||
_mock_chat_response("not valid {{{ json"),
|
||||
_mock_chat_response(json.dumps({
|
||||
"title_zh": "测试标题",
|
||||
"content_zh": "测试正文",
|
||||
"events": [],
|
||||
})),
|
||||
]
|
||||
|
||||
result = translate_and_extract(
|
||||
client=mock_client,
|
||||
config=config,
|
||||
article=article,
|
||||
max_attempts=3,
|
||||
)
|
||||
|
||||
assert result.attempts == 2 # 第二次成功
|
||||
assert result.title_zh == "测试标题"
|
||||
|
||||
|
||||
def _mock_chat_response(content: str) -> MagicMock:
|
||||
"""Helper：构造 mock OpenAI chat completion 响应。"""
|
||||
resp = MagicMock()
|
||||
resp.choices = [MagicMock()]
|
||||
resp.choices[0].message.content = content
|
||||
resp.usage = MagicMock()
|
||||
resp.usage.prompt_tokens = 100
|
||||
resp.usage.completion_tokens = 50
|
||||
return resp
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# INTERNATIONAL_EVENT_TYPES
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestEventTypes:
|
||||
"""事件类型常量测试。"""
|
||||
|
||||
def test_has_expected_types(self):
|
||||
assert "财报披露" in INTERNATIONAL_EVENT_TYPES
|
||||
assert "并购收购" in INTERNATIONAL_EVENT_TYPES
|
||||
assert "央行决议" in INTERNATIONAL_EVENT_TYPES
|
||||
assert "地缘政治" in INTERNATIONAL_EVENT_TYPES
|
||||
assert len(INTERNATIONAL_EVENT_TYPES) >= 10
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# LLMCallError
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class TestLLMCallError:
|
||||
"""LLMCallError 异常测试。"""
|
||||
|
||||
def test_basic(self):
|
||||
err = LLMCallError("测试错误", attempts=3)
|
||||
assert err.reason == "测试错误"
|
||||
assert err.attempts == 3
|
||||
assert str(err) == "测试错误"
|
||||
|
||||
def test_default_attempts(self):
|
||||
err = LLMCallError("错误")
|
||||
assert err.attempts == 0
|
||||
Reference in New Issue
Block a user