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

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
StateManager extra tests
"""

import json
import sys
import tempfile
from pathlib import Path

import pytest

from data_modules.state_manager import StateManager, EntityState
from data_modules.index_manager import IndexManager, EntityMeta


@pytest.fixture
def temp_project(tmp_path):
    from data_modules.config import DataModulesConfig
    cfg = DataModulesConfig.from_project_root(tmp_path)
    cfg.ensure_dirs()
    return cfg


def test_ensure_state_schema_and_progress(temp_project):
    # relationships as list should be migrated to structured_relationships
    state = {
        "relationships": [
            {"from_entity": "a", "to_entity": "b", "type": "师徒", "chapter": 1}
        ],
        "progress": {"current_chapter": "2", "total_words": "10"},
    }
    temp_project.state_file.write_text(json.dumps(state, ensure_ascii=False), encoding="utf-8")

    manager = StateManager(temp_project, enable_sqlite_sync=False)
    assert isinstance(manager._state.get("relationships"), dict)
    assert isinstance(manager._state.get("structured_relationships"), list)
    assert int(manager.get_current_chapter()) == 2

    manager.update_progress(3)
    assert manager.get_current_chapter() == 3


def test_add_update_entities_and_alias(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)

    entity = EntityState(id="xiaoyan", name="萧炎", type="角色", tier="核心", aliases=["炎帝"])
    assert manager.add_entity(entity) is True
    assert manager.add_entity(entity) is False

    manager.update_entity("xiaoyan", {"current": {"realm": "斗师"}})
    updated = manager.get_entity("xiaoyan")
    assert updated["current"]["realm"] == "斗师"

    assert manager.get_entity_type("xiaoyan") == "角色"
    assert manager.get_entity_type("missing") is None

    assert "xiaoyan" in manager.get_all_entities()
    assert "xiaoyan" in manager.get_entities_by_type("角色")
    assert "xiaoyan" in manager.get_entities_by_tier("核心")

    # unknown type update
    assert manager.update_entity("missing", {"current": {"realm": "斗者"}}, "角色") is False


def test_update_entity_appearance_and_relationships(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)
    manager.add_entity(EntityState(id="xiaoyan", name="萧炎", type="角色"))

    manager.update_entity_appearance("xiaoyan", 5, "角色")
    entity = manager.get_entity("xiaoyan")
    assert entity.get("first_appearance") == 5
    assert entity.get("last_appearance") == 5

    # unknown entity should no-op
    manager.update_entity_appearance("missing", 3, "角色")

    manager.add_relationship("xiaoyan", "yaolao", "师徒", "收徒", 1)
    rels = manager.get_relationships("xiaoyan")
    assert len(rels) == 1


def test_disambiguation_and_save_state(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)
    warnings = manager._record_disambiguation(
        1,
        [
            {
                "mention": "宗主",
                "candidates": ["zongzhu", "lintian"],
                "suggested": "zongzhu",
                "confidence": 0.4,
            },
            {
                "mention": "萧炎",
                "candidates": [{"type": "角色", "id": "xiaoyan"}],
                "suggested": "xiaoyan",
                "confidence": 0.6,
            },
        ],
    )
    assert any("需人工确认" in w for w in warnings)
    assert any("消歧警告" in w for w in warnings)

    manager.save_state()
    assert temp_project.state_file.exists()


def test_save_state_no_pending(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)
    manager.save_state()
    assert not temp_project.state_file.exists()


def test_save_state_with_sqlite_sync_and_protagonist(temp_project):
    manager = StateManager(temp_project)
    manager.add_entity(EntityState(id="xiaoyan", name="萧炎", type="角色", tier="核心"))
    manager.update_entity("xiaoyan", {"current": {"realm": "斗师", "location": "天云宗"}})
    manager.update_progress(10, words=500)
    manager.save_state()

    state = json.loads(temp_project.state_file.read_text(encoding="utf-8"))
    assert state.get("_migrated_to_sqlite") is True
    assert state.get("progress", {}).get("current_chapter") == 10

    # 标记为主角并同步
    idx = IndexManager(temp_project)
    idx.upsert_entity(
        EntityMeta(
            id="xiaoyan",
            type="角色",
            canonical_name="萧炎",
            tier="核心",
            current={"realm": "斗王", "location": "天云宗"},
            first_appearance=1,
            last_appearance=10,
            is_protagonist=True,
        ),
        update_metadata=True,
    )
    manager.sync_protagonist_from_entity()
    assert manager._state.get("protagonist_state", {}).get("power", {}).get("realm") == "斗王"

    manager._state["protagonist_state"] = {
        "power": {"realm": "斗皇", "layer": 2},
        "location": {"current": "中州"},
    }
    manager._state.setdefault("entities_v3", {"角色": {}})
    manager._state["entities_v3"]["角色"]["xiaoyan"] = {
        "canonical_name": "萧炎",
        "tier": "核心",
        "desc": "",
        "current": {"realm": "斗王", "location": "天云宗"},
        "first_appearance": 1,
        "last_appearance": 10,
        "history": [],
    }
    manager.sync_protagonist_to_entity("xiaoyan")
    manager.save_state()
    updated = idx.get_entity("xiaoyan")
    assert updated["current_json"]["realm"] == "斗皇"

    # export context
    exported = manager.export_for_context()
    assert exported.get("alias_index") == {}


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

    result = {
        "entities_appeared": [
            {"id": "xiaoyan", "type": "角色", "mentions": ["萧炎"], "confidence": 0.9}
        ],
        "entities_new": [
            {
                "suggested_id": "yaolao",
                "name": "药老",
                "type": "角色",
                "tier": "重要",
                "mentions": ["药老"],
                "aliases": ["药老先生"],
            }
        ],
        "state_changes": [
            {"entity_id": "xiaoyan", "field": "realm", "old": "斗者", "new": "斗师", "reason": "突破"}
        ],
        "relationships_new": [
            {"from": "xiaoyan", "to": "yaolao", "type": "师徒", "description": "收徒"}
        ],
        "uncertain": [
            {"mention": "宗主", "candidates": ["zongzhu", "lintian"], "suggested": "zongzhu", "confidence": 0.2},
            {
                "mention": "萧炎",
                "candidates": [{"type": "角色", "id": "xiaoyan"}],
                "suggested": "xiaoyan",
                "confidence": 0.8,
                "adopted": True,
            },
        ],
        "chapter_meta": {"hook": "test", "end": "ok"},
    }
    warnings = manager.process_chapter_result(12, result)
    assert any("需人工确认" in w for w in warnings)
    assert any("消歧警告" in w for w in warnings)

    manager.save_state()

    idx = IndexManager(temp_project)
    assert idx.get_entity("yaolao") is not None
    assert idx.get_relationship_between("xiaoyan", "yaolao")
    assert idx.get_entity_state_changes("xiaoyan")

    by_type = manager.get_entities_by_type("角色")
    by_tier = manager.get_entities_by_tier("核心")
    assert "xiaoyan" in by_type
    assert "xiaoyan" in by_tier


def test_export_context_and_protagonist_alias(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)
    manager.add_entity(EntityState(id="xiaoyan", name="萧炎", type="角色", tier="核心"))
    manager._state["disambiguation_warnings"] = [{"chapter": 1, "mention": "萧炎"}]
    manager._state["disambiguation_pending"] = [{"chapter": 2, "mention": "宗主"}]

    exported = manager.export_for_context()
    assert "xiaoyan" in exported.get("entities", {})
    assert exported["disambiguation"]["warnings"]
    assert exported["disambiguation"]["pending"]

    manager_sql = StateManager(temp_project)
    idx = IndexManager(temp_project)
    idx.upsert_entity(
        EntityMeta(
            id="xiaoyan",
            type="角色",
            canonical_name="萧炎",
            tier="核心",
            current={},
            first_appearance=1,
            last_appearance=1,
            is_protagonist=False,
        ),
        update_metadata=True,
    )
    idx.register_alias("小炎子", "xiaoyan", "角色")
    manager_sql._state["protagonist_state"] = {"name": "小炎子"}
    assert manager_sql.get_protagonist_entity_id() == "xiaoyan"

    idx.upsert_entity(
        EntityMeta(
            id="xiaoyan",
            type="角色",
            canonical_name="萧炎",
            tier="核心",
            current={},
            first_appearance=1,
            last_appearance=1,
            is_protagonist=True,
        ),
        update_metadata=True,
    )
    assert manager_sql.get_protagonist_entity_id() == "xiaoyan"


def test_sqlite_metadata_update_and_alias_sync(temp_project):
    manager = StateManager(temp_project)
    idx = IndexManager(temp_project)
    idx.upsert_entity(
        EntityMeta(
            id="xiaoyan",
            type="角色",
            canonical_name="萧炎",
            tier="核心",
            current={"realm": "斗者"},
            first_appearance=1,
            last_appearance=1,
            is_protagonist=False,
        )
    )

    manager._state.setdefault("entities_v3", {"角色": {}})
    manager._state["entities_v3"]["角色"]["xiaoyan"] = {
        "canonical_name": "萧炎",
        "tier": "核心",
        "desc": "",
        "current": {"realm": "斗者"},
        "first_appearance": 1,
        "last_appearance": 1,
        "history": [],
    }

    manager.update_entity(
        "xiaoyan",
        {"canonical_name": "萧炎·新", "tier": "重要", "current": {"realm": "斗王"}},
        "角色",
    )
    manager.update_entity("xiaoyan", {"location": "中州"}, "角色")
    manager.update_entity_appearance("xiaoyan", 2, "角色")
    manager._pending_alias_entries["小炎子"] = [{"type": "角色", "id": "xiaoyan"}]

    manager.save_state()

    updated = idx.get_entity("xiaoyan")
    assert updated["canonical_name"] == "萧炎·新"
    assert updated["current_json"]["realm"] == "斗王"
    assert updated["current_json"]["location"] == "中州"
    assert updated["last_appearance"] == 2

    aliases = idx.get_entity_aliases("xiaoyan")
    assert "萧炎·新" in aliases
    assert "小炎子" in aliases


def test_ensure_state_schema_invalid_inputs(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)
    schema = manager._ensure_state_schema("bad")
    assert isinstance(schema, dict)

    schema2 = manager._ensure_state_schema({
        "progress": "bad",
        "relationships": "bad",
        "disambiguation_warnings": "bad",
        "disambiguation_pending": "bad",
    })
    assert isinstance(schema2["progress"], dict)
    assert isinstance(schema2["relationships"], dict)
    assert isinstance(schema2["disambiguation_warnings"], list)
    assert isinstance(schema2["disambiguation_pending"], list)


def test_save_state_preserves_sqlite_pending_on_sync_failure(temp_project):
    manager = StateManager(temp_project)

    manager.add_entity(EntityState(id="e1", name="测试角色", type="角色", first_appearance=1, last_appearance=1))
    manager.update_entity("e1", {"current": {"realm": "炼气"}}, "角色")

    class _BrokenSQLManager:
        def process_chapter_entities(self, **kwargs):
            raise RuntimeError("boom")

    manager._sql_state_manager = _BrokenSQLManager()
    manager._pending_sqlite_data["chapter"] = 1

    manager.save_state()

    state = json.loads(temp_project.state_file.read_text(encoding="utf-8"))
    assert state.get("_migrated_to_sqlite") is True

    # SQLite 同步失败后，SQLite 相关 pending 不应被清空，便于后续重试
    assert manager._pending_entity_patches
    assert manager._pending_sqlite_data.get("chapter") == 1


def test_save_state_progress_and_disambiguation_merge(temp_project):
    state = {
        "progress": {"current_chapter": "bad", "total_words": "bad"},
        "disambiguation_warnings": "bad",
        "disambiguation_pending": "bad",
    }
    temp_project.state_file.write_text(json.dumps(state, ensure_ascii=False), encoding="utf-8")

    manager = StateManager(temp_project, enable_sqlite_sync=False)
    manager.config.max_disambiguation_warnings = 1
    manager.config.max_disambiguation_pending = 1
    manager._pending_progress_chapter = 5
    manager._pending_progress_words_delta = 10
    manager._pending_disambiguation_warnings = [
        {"chapter": 1, "mention": "a", "chosen_id": "x", "confidence": 0.5},
        {"chapter": 1, "mention": "a", "chosen_id": "x", "confidence": 0.5},
        "bad",
    ]
    manager._pending_disambiguation_pending = [
        {"chapter": 2, "mention": "b", "suggested_id": "y", "confidence": 0.4},
        {"chapter": 2, "mention": "b", "suggested_id": "y", "confidence": 0.4},
        "bad",
    ]
    manager.save_state()

    saved = json.loads(temp_project.state_file.read_text(encoding="utf-8"))
    assert saved["progress"]["current_chapter"] == 5
    assert saved["progress"]["total_words"] == 10
    assert len(saved["disambiguation_warnings"]) == 1
    assert len(saved["disambiguation_pending"]) == 1


def test_sync_to_sqlite_exceptions_and_no_sql_manager(temp_project, monkeypatch):
    manager = StateManager(temp_project)
    manager._pending_progress_chapter = 1
    manager._pending_sqlite_data["chapter"] = 1
    manager._pending_alias_entries["alias"] = [{"type": "角色", "id": "xiaoyan"}]

    def boom(*args, **kwargs):
        raise RuntimeError("boom")

    monkeypatch.setattr(manager._sql_state_manager, "process_chapter_entities", boom)
    monkeypatch.setattr(manager._sql_state_manager, "register_alias", boom)

    manager.save_state()

    manager_no_sql = StateManager(temp_project, enable_sqlite_sync=False)
    manager_no_sql._sync_pending_patches_to_sqlite()


def test_entity_fallbacks_and_updates(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)

    manager.add_entity(EntityState(id="hero", name="主角", type="未知", tier="核心"))
    manager.add_entity(EntityState(id="place", name="乌坦城", type="地点", tier="重要"))

    assert manager.get_entity("hero", "角色")["canonical_name"] == "主角"
    assert manager.get_entity("place")["canonical_name"] == "乌坦城"
    assert manager.get_entity_type("place") == "地点"

    assert "hero" in manager.get_entities_by_type("角色")
    assert "hero" in manager.get_entities_by_tier("核心")
    assert "hero" in manager.get_all_entities()

    assert manager.update_entity("missing", {"current": {"a": 1}}) is False

    manager.update_entity("hero", {"attributes": {"hp": 1}}, "角色")
    manager._state["entities_v3"]["角色"]["hero"].pop("current", None)
    manager.update_entity("hero", {"current": {"mp": 2}}, "角色")
    manager.update_entity("hero", {"tier": "重要"}, "角色")

    manager._state["entities_v3"] = "bad"
    manager.update_entity_appearance("hero", 1, "角色")
    manager._state["entities_v3"]["角色"]["hero"] = {"first_appearance": 0, "last_appearance": 0}
    manager.update_entity_appearance("hero", 1, "角色")
    manager.update_entity_appearance("hero", 2, "角色")


def test_register_alias_internal_and_get_all_entities_sqlite(temp_project):
    manager = StateManager(temp_project)
    manager._register_alias_internal("xiaoyan", "角色", "")
    manager._register_alias_internal("xiaoyan", "角色", "萧炎")

    idx = IndexManager(temp_project)
    idx.upsert_entity(
        EntityMeta(
            id="xiaoyan",
            type="角色",
            canonical_name="萧炎",
            tier="核心",
            current={},
            first_appearance=1,
            last_appearance=1,
            is_protagonist=False,
        )
    )
    all_entities = manager.get_all_entities()
    assert "xiaoyan" in all_entities


def test_record_disambiguation_and_process_chapter_existing(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)
    warnings = manager._record_disambiguation(
        1,
        [
            "bad",
            {"mention": "", "confidence": 0.1},
            {"mention": "宗主", "confidence": "bad", "adopted": "zongzhu"},
        ],
    )
    assert warnings

    manager.add_entity(EntityState(id="xiaoyan", name="萧炎", type="角色"))
    warnings = manager.process_chapter_result(2, {"entities_new": [{"id": "xiaoyan", "name": "萧炎"}]})
    assert any("实体已存在" in w for w in warnings)


def test_sync_protagonist_from_string_and_empty_updates(temp_project):
    manager = StateManager(temp_project, enable_sqlite_sync=False)
    manager._state.setdefault("entities_v3", {"角色": {}})
    manager._state["entities_v3"]["角色"]["bad"] = {
        "current": None,
        "current_json": "not-json",
    }
    manager._state["entities_v3"]["角色"]["hero"] = {
        "current": None,
        "current_json": json.dumps({"realm": "斗师", "layer": 2, "location": "乌坦城", "last_chapter": 3}),
    }
    manager.sync_protagonist_from_entity("bad")
    manager.sync_protagonist_from_entity("hero")
    assert manager._state["protagonist_state"]["power"]["realm"] == "斗师"

    manager._state["protagonist_state"] = {}
    manager.sync_protagonist_to_entity()


def test_state_manager_cli_commands(temp_project, monkeypatch, capsys):
    idx = IndexManager(temp_project)
    idx.upsert_entity(
        EntityMeta(
            id="xiaoyan",
            type="角色",
            canonical_name="萧炎",
            tier="核心",
            current={},
            first_appearance=1,
            last_appearance=1,
            is_protagonist=False,
        )
    )

    def run_cli(args):
        monkeypatch.setattr(sys, "argv", args)
        from data_modules import state_manager as sm

        sm.main()
        out = capsys.readouterr().out
        return json.loads(out)

    out = run_cli(["state_manager", "--project-root", str(temp_project.project_root), "get-progress"])
    assert out["status"] == "success"
    assert "current_chapter" in out.get("data", {})

    out = run_cli(["state_manager", "--project-root", str(temp_project.project_root), "get-entity", "--id", "missing"])
    assert out["status"] == "error"

    out = run_cli(["state_manager", "--project-root", str(temp_project.project_root), "get-entity", "--id", "xiaoyan"])
    assert out["status"] == "success"
    assert out["data"].get("id") == "xiaoyan"

    out = run_cli(["state_manager", "--project-root", str(temp_project.project_root), "list-entities", "--type", "角色"])
    assert out["status"] == "success"
    assert any(e.get("id") == "xiaoyan" for e in out.get("data", []))

    out = run_cli(["state_manager", "--project-root", str(temp_project.project_root), "list-entities", "--tier", "核心"])
    assert out["status"] == "success"
    assert any(e.get("id") == "xiaoyan" for e in out.get("data", []))

    payload = json.dumps({"entities_appeared": [], "entities_new": [], "state_changes": [], "relationships_new": []})
    out = run_cli([
        "state_manager",
        "--project-root",
        str(temp_project.project_root),
        "process-chapter",
        "--chapter",
        "1",
        "--data",
        payload,
    ])
    assert out["status"] == "success"


def test_save_state_timeout(monkeypatch, temp_project):
    import filelock
    from data_modules import state_manager as sm

    manager = StateManager(temp_project, enable_sqlite_sync=False)
    manager.update_progress(1)

    class FakeLock:
        def __init__(self, *args, **kwargs):
            pass

        def __enter__(self):
            raise filelock.Timeout("timeout")

        def __exit__(self, exc_type, exc, tb):
            return False

    monkeypatch.setattr(sm.filelock, "FileLock", FakeLock)
    with pytest.raises(RuntimeError):
        manager.save_state()