#!/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()