"""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")