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

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

import sys

import pytest

from data_modules.entity_linker import EntityLinker, main as linker_main
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_process_extraction_and_register_new_entities(temp_project):
    linker = EntityLinker(temp_project)
    idx = IndexManager(temp_project)
    idx.upsert_entity(
        EntityMeta(
            id="xiaoyan",
            type="角色",
            canonical_name="萧炎",
            current={},
            first_appearance=1,
            last_appearance=1,
        )
    )

    results, warnings = linker.process_extraction_result(
        [
            {
                "mention": "萧炎",
                "candidates": ["xiaoyan"],
                "suggested": "xiaoyan",
                "confidence": 0.7,
            },
            {
                "mention": "宗主",
                "candidates": ["zongzhu"],
                "suggested": "zongzhu",
                "confidence": 0.4,
            },
        ]
    )

    assert len(results) == 2
    assert len(warnings) == 2

    registered = linker.register_new_entities(
        [
            {
                "suggested_id": "hongyi",
                "name": "红衣女子",
                "type": "角色",
                "mentions": ["红衣", "女子"],
            }
        ]
    )
    assert registered == ["hongyi"]
    aliases = idx.get_entity_aliases("hongyi")
    assert "红衣女子" in aliases


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

    def run_cli(args):
        monkeypatch.setattr(sys, "argv", ["entity_linker"] + args)
        linker_main()

    root = str(temp_project.project_root)

    run_cli(["--project-root", root, "register-alias", "--entity", "xiaoyan", "--alias", "炎帝"])
    run_cli(["--project-root", root, "lookup", "--mention", "炎帝"])
    run_cli(["--project-root", root, "lookup", "--mention", "不存在"])
    run_cli(["--project-root", root, "lookup-all", "--mention", "炎帝"])
    run_cli(["--project-root", root, "list-aliases", "--entity", "xiaoyan"])

    capsys.readouterr()