webnovel-writer/.claude/scripts/data_modules/tests/test_data_modules.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Data Modules 单元测试
"""

import pytest
import asyncio
import json
import tempfile
from pathlib import Path

from data_modules import (
    DataModulesConfig,
    EntityLinker,
    StateManager,
    IndexManager,
    RAGAdapter,
    StyleSampler,
    EntityState,
    ChapterMeta,
    SceneMeta,
    StyleSample,
)


@pytest.fixture
def temp_project():
    """创建临时项目目录"""
    with tempfile.TemporaryDirectory() as tmpdir:
        config = DataModulesConfig.from_project_root(tmpdir)
        config.ensure_dirs()
        yield config


class TestEntityLinker:
    """实体链接器测试"""

    def test_register_and_lookup_alias(self, temp_project):
        linker = EntityLinker(temp_project)

        # 注册别名
        assert linker.register_alias("xiaoyan", "萧炎")
        assert linker.register_alias("xiaoyan", "小炎子")

        # 查找
        assert linker.lookup_alias("萧炎") == "xiaoyan"
        assert linker.lookup_alias("小炎子") == "xiaoyan"
        assert linker.lookup_alias("不存在") is None

    def test_alias_one_to_many(self, temp_project):
        """v5.0: 同一别名可映射多个实体（一对多）"""
        linker = EntityLinker(temp_project)

        linker.register_alias("xiaoyan", "萧炎", "角色")
        # v5.0: 同一别名可绑定不同实体（一对多）
        assert linker.register_alias("other_person", "萧炎", "角色")

        # 查找所有匹配
        entries = linker.lookup_alias_all("萧炎")
        assert len(entries) == 2

    def test_get_all_aliases(self, temp_project):
        linker = EntityLinker(temp_project)

        linker.register_alias("xiaoyan", "萧炎")
        linker.register_alias("xiaoyan", "小炎子")
        linker.register_alias("xiaoyan", "炎哥")

        aliases = linker.get_all_aliases("xiaoyan")
        assert len(aliases) == 3
        assert "萧炎" in aliases

    def test_confidence_evaluation(self, temp_project):
        linker = EntityLinker(temp_project)

        # 高置信度
        action, adopt, warning = linker.evaluate_confidence(0.9)
        assert action == "auto"
        assert adopt is True
        assert warning is None

        # 中置信度
        action, adopt, warning = linker.evaluate_confidence(0.6)
        assert action == "warn"
        assert adopt is True
        assert warning is not None

        # 低置信度
        action, adopt, warning = linker.evaluate_confidence(0.3)
        assert action == "manual"
        assert adopt is False

    def test_process_uncertain(self, temp_project):
        linker = EntityLinker(temp_project)

        result = linker.process_uncertain(
            mention="那位前辈",
            candidates=["yaolao", "elder_zhang"],
            suggested="yaolao",
            confidence=0.7
        )

        assert result.mention == "那位前辈"
        assert result.entity_id == "yaolao"
        assert result.adopted is True
        assert result.warning is not None


class TestStateManager:
    """状态管理器测试"""

    def test_add_and_get_entity(self, temp_project):
        manager = StateManager(temp_project)

        entity = EntityState(
            id="xiaoyan",
            name="萧炎",
            type="角色",
            tier="核心"
        )
        assert manager.add_entity(entity)

        # 获取实体
        result = manager.get_entity("xiaoyan")
        assert result is not None
        assert result["canonical_name"] == "萧炎"

    def test_update_entity(self, temp_project):
        manager = StateManager(temp_project)

        entity = EntityState(id="xiaoyan", name="萧炎", type="角色")
        manager.add_entity(entity)

        # 更新属性 (v5.0: attributes 存在 current 字段)
        manager.update_entity("xiaoyan", {"current": {"realm": "斗师"}})

        result = manager.get_entity("xiaoyan")
        assert result["current"]["realm"] == "斗师"

    def test_record_state_change(self, temp_project):
        manager = StateManager(temp_project)

        entity = EntityState(id="xiaoyan", name="萧炎", type="角色")
        manager.add_entity(entity)

        manager.record_state_change(
            entity_id="xiaoyan",
            field="realm",
            old_value="斗者",
            new_value="斗师",
            reason="突破",
            chapter=100
        )

        changes = manager.get_state_changes("xiaoyan")
        assert len(changes) == 1
        assert changes[0]["new_value"] == "斗师"

    def test_add_relationship(self, temp_project):
        manager = StateManager(temp_project)

        manager.add_relationship(
            from_entity="xiaoyan",
            to_entity="yaolao",
            rel_type="师徒",
            description="药老收萧炎为徒",
            chapter=10
        )

        rels = manager.get_relationships("xiaoyan")
        assert len(rels) == 1
        assert rels[0]["type"] == "师徒"

    def test_process_chapter_result(self, temp_project):
        manager = StateManager(temp_project)

        result = {
            "entities_appeared": [
                {"id": "xiaoyan", "mentions": ["萧炎", "他"]}
            ],
            "entities_new": [
                {"suggested_id": "hongyi_girl", "name": "红衣女子", "type": "角色", "tier": "装饰"}
            ],
            "state_changes": [
                {"entity_id": "xiaoyan", "field": "realm", "old": "斗者", "new": "斗师", "reason": "突破"}
            ],
            "relationships_new": [
                {"from": "xiaoyan", "to": "hongyi_girl", "type": "相识", "description": "初次见面"}
            ]
        }

        # 先添加萧炎
        manager.add_entity(EntityState(id="xiaoyan", name="萧炎", type="角色"))

        warnings = manager.process_chapter_result(100, result)

        # 验证新实体被添加
        assert manager.get_entity("hongyi_girl") is not None

        # 验证状态变化
        changes = manager.get_state_changes("xiaoyan")
        assert len(changes) == 1

        # 验证进度更新
        assert manager.get_current_chapter() == 100

    def test_save_state_with_init_project_schema(self, temp_project):
        """回归：init_project 生成的 state.json 无 meta 字段，StateManager 仍应可写入。"""
        # 模拟 init_project.py 生成的 v5.0 state.json 形状（包含 entities_v3/alias_index）
        init_state = {
            "project_info": {"title": "测试书名", "genre": "修仙/玄幻", "created_at": "2026-01-01"},
            "progress": {"current_chapter": 0, "total_words": 0, "last_updated": "2026-01-01 00:00:00"},
            "protagonist_state": {"name": "测试主角"},
            "relationships": {},
            "world_settings": {"power_system": [], "factions": [], "locations": []},
            "plot_threads": {"active_threads": [], "foreshadowing": []},
            "review_checkpoints": [],
            "strand_tracker": {"current_dominant": "quest", "history": []},
            "entities_v3": {"角色": {}, "地点": {}, "物品": {}, "势力": {}, "招式": {}},
            "alias_index": {},
        }
        temp_project.state_file.write_text(json.dumps(init_state, ensure_ascii=False, indent=2), encoding="utf-8")

        manager = StateManager(temp_project)
        manager.update_progress(5, words=100)
        manager.save_state()

        saved = json.loads(temp_project.state_file.read_text(encoding="utf-8"))
        assert "meta" not in saved
        assert saved["progress"]["current_chapter"] == 5
        assert saved["progress"]["total_words"] == 100
        assert isinstance(saved.get("entities_v3"), dict)
        assert isinstance(saved.get("alias_index"), dict)

    def test_save_state_preserves_unrelated_fields(self, temp_project):
        """回归：仅写入增量，不应覆盖/丢失其他模块维护的字段。"""
        init_state = {
            "project_info": {"title": "测试书名", "genre": "修仙/玄幻", "created_at": "2026-01-01"},
            "progress": {"current_chapter": 10, "total_words": 1000, "last_updated": "2026-01-01 00:00:00"},
            "protagonist_state": {"name": "测试主角"},
            "relationships": {"allies": ["药老"], "enemies": []},
            "world_settings": {"power_system": [], "factions": [], "locations": []},
            "plot_threads": {"active_threads": [{"id": "t1", "title": "主线"}], "foreshadowing": []},
            "review_checkpoints": [],
            "strand_tracker": {"current_dominant": "quest", "history": []},
            "entities_v3": {"角色": {}, "地点": {}, "物品": {}, "势力": {}, "招式": {}},
            "alias_index": {},
            "custom_field": {"keep": True},
        }
        temp_project.state_file.write_text(json.dumps(init_state, ensure_ascii=False, indent=2), encoding="utf-8")

        manager = StateManager(temp_project)
        manager.add_entity(EntityState(id="xiaoyan", name="萧炎", type="角色", tier="核心"))
        manager.save_state()

        saved = json.loads(temp_project.state_file.read_text(encoding="utf-8"))
        assert saved.get("custom_field", {}).get("keep") is True
        assert saved.get("plot_threads", {}).get("active_threads", [])[0].get("id") == "t1"
        assert isinstance(saved.get("relationships"), dict)

    def test_disambiguation_feedback_persisted(self, temp_project):
        """回归：中/低置信度消歧必须对 Writer 可见（写入 state.json）。"""
        manager = StateManager(temp_project)

        result = {
            "entities_appeared": [],
            "entities_new": [],
            "state_changes": [],
            "relationships_new": [],
            "uncertain": [
                {
                    "mention": "那位前辈",
                    "context": "那位前辈看了他一眼",
                    "candidates": [{"type": "角色", "id": "yaolao"}, {"type": "角色", "id": "elder_zhang"}],
                    "suggested": "yaolao",
                    "confidence": 0.6,
                },
                {
                    "mention": "宗主",
                    "context": "宗主出现在血煞秘境",
                    "candidates": ["xueshazonzhu", "lintian"],
                    "suggested": "xueshazonzhu",
                    "confidence": 0.4,
                },
            ],
        }

        warnings = manager.process_chapter_result(100, result)
        manager.save_state()

        state = json.loads(temp_project.state_file.read_text(encoding="utf-8"))
        assert isinstance(state.get("disambiguation_warnings"), list)
        assert isinstance(state.get("disambiguation_pending"), list)

        assert len(state["disambiguation_warnings"]) == 1
        assert len(state["disambiguation_pending"]) == 1

        warn = state["disambiguation_warnings"][0]
        assert warn.get("chapter") == 100
        assert warn.get("mention") == "那位前辈"
        assert warn.get("chosen_id") == "yaolao"

        pending = state["disambiguation_pending"][0]
        assert pending.get("chapter") == 100
        assert pending.get("mention") == "宗主"

        # 返回值也应包含可见警告，便于 CLI/日志透出
        assert any("消歧警告" in w for w in warnings)
        assert any("需人工确认" in w for w in warnings)


class TestIndexManager:
    """索引管理器测试"""

    def test_add_and_get_chapter(self, temp_project):
        manager = IndexManager(temp_project)

        meta = ChapterMeta(
            chapter=100,
            title="突破",
            location="天云宗",
            word_count=3500,
            characters=["xiaoyan", "yaolao"]
        )
        manager.add_chapter(meta)

        result = manager.get_chapter(100)
        assert result is not None
        assert result["title"] == "突破"
        assert "xiaoyan" in result["characters"]

    def test_add_scenes(self, temp_project):
        manager = IndexManager(temp_project)

        scenes = [
            SceneMeta(chapter=100, scene_index=1, start_line=1, end_line=50,
                     location="天云宗·闭关室", summary="萧炎闭关突破", characters=["xiaoyan"]),
            SceneMeta(chapter=100, scene_index=2, start_line=51, end_line=100,
                     location="天云宗·演武场", summary="展示实力", characters=["xiaoyan", "lintian"])
        ]
        manager.add_scenes(100, scenes)

        result = manager.get_scenes(100)
        assert len(result) == 2
        assert result[0]["location"] == "天云宗·闭关室"

    def test_record_appearance(self, temp_project):
        manager = IndexManager(temp_project)

        manager.record_appearance("xiaoyan", 100, ["萧炎", "他"], 0.95)
        manager.record_appearance("yaolao", 100, ["药老"], 0.92)

        appearances = manager.get_chapter_appearances(100)
        assert len(appearances) == 2

        entity_history = manager.get_entity_appearances("xiaoyan")
        assert len(entity_history) == 1

    def test_search_scenes_by_location(self, temp_project):
        manager = IndexManager(temp_project)

        scenes = [
            SceneMeta(chapter=100, scene_index=1, start_line=1, end_line=50,
                     location="天云宗·闭关室", summary="闭关", characters=[]),
            SceneMeta(chapter=101, scene_index=1, start_line=1, end_line=50,
                     location="天云宗·大殿", summary="议事", characters=[])
        ]
        manager.add_scenes(100, scenes[:1])
        manager.add_scenes(101, scenes[1:])

        results = manager.search_scenes_by_location("天云宗")
        assert len(results) == 2

    def test_get_stats(self, temp_project):
        manager = IndexManager(temp_project)

        manager.add_chapter(ChapterMeta(chapter=1, title="", location="", word_count=1000, characters=[]))
        manager.add_scenes(1, [SceneMeta(chapter=1, scene_index=1, start_line=1, end_line=50,
                                        location="", summary="", characters=[])])
        manager.record_appearance("xiaoyan", 1, [], 1.0)

        stats = manager.get_stats()
        assert stats["chapters"] == 1
        assert stats["scenes"] == 1
        assert stats["entities"] == 1


class TestStyleSampler:
    """风格样本测试"""

    def test_add_and_get_sample(self, temp_project):
        sampler = StyleSampler(temp_project)

        sample = StyleSample(
            id="ch100_s1",
            chapter=100,
            scene_type="战斗",
            content="萧炎一拳轰出...",
            score=0.85,
            tags=["战斗", "激烈"]
        )
        assert sampler.add_sample(sample)

        results = sampler.get_samples_by_type("战斗")
        assert len(results) == 1
        assert results[0].id == "ch100_s1"

    def test_extract_candidates(self, temp_project):
        sampler = StyleSampler(temp_project)

        scenes = [
            {"index": 1, "summary": "战斗场景", "content": "萧炎一拳轰出，斗气如虹，直接将对手击退三丈，周围的空气都被震得嗡嗡作响..." + "a" * 200}
        ]

        # 低分不提取
        candidates = sampler.extract_candidates(100, "", 70, scenes)
        assert len(candidates) == 0

        # 高分提取
        candidates = sampler.extract_candidates(100, "", 85, scenes)
        assert len(candidates) == 1
        assert candidates[0].scene_type == "战斗"

    def test_select_samples_for_chapter(self, temp_project):
        sampler = StyleSampler(temp_project)

        # 添加一些样本
        for i in range(3):
            sampler.add_sample(StyleSample(
                id=f"battle_{i}",
                chapter=i,
                scene_type="战斗",
                content=f"战斗内容 {i}",
                score=0.9,
                tags=[]
            ))

        samples = sampler.select_samples_for_chapter("本章有一场激烈的战斗")
        assert len(samples) <= 3
        assert all(s.scene_type == "战斗" for s in samples)


class TestRAGAdapter:
    """RAG 适配器测试（不包含 API 调用）"""

    def test_bm25_search(self, temp_project):
        adapter = RAGAdapter(temp_project)

        # 手动插入一些测试数据
        with adapter._get_conn() as conn:
            cursor = conn.cursor()

            # 插入向量记录（空向量，只测试 BM25）
            cursor.execute("""
                INSERT INTO vectors (chunk_id, chapter, scene_index, content, embedding)
                VALUES (?, ?, ?, ?, ?)
            """, ("ch1_s1", 1, 1, "萧炎在天云宗修炼斗气", b""))

            cursor.execute("""
                INSERT INTO vectors (chunk_id, chapter, scene_index, content, embedding)
                VALUES (?, ?, ?, ?, ?)
            """, ("ch1_s2", 1, 2, "药老传授炼药技巧", b""))

            conn.commit()

            # 更新 BM25 索引
            adapter._update_bm25_index(cursor, "ch1_s1", "萧炎在天云宗修炼斗气")
            adapter._update_bm25_index(cursor, "ch1_s2", "药老传授炼药技巧")
            conn.commit()

        # BM25 搜索
        results = adapter.bm25_search("萧炎修炼", top_k=5)
        assert len(results) >= 1
        assert results[0].chunk_id == "ch1_s1"

    def test_tokenize(self, temp_project):
        adapter = RAGAdapter(temp_project)

        tokens = adapter._tokenize("萧炎hello世界world")
        assert "萧" in tokens
        assert "炎" in tokens
        assert "hello" in tokens
        assert "world" in tokens


if __name__ == "__main__":
    pytest.main([__file__, "-v"])