vent_agent_2.0/tests/test_extensions.py

"""业务扩展发现、注册、组装和启动失败路径。"""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import sys
from types import SimpleNamespace

import pytest

from extensions import ExtensionRegistry, define_tool
from extensions import registry as extensions
from tools import registry as tools


@define_tool(name="airflow", group="通风", label="计算风量")
def airflow_impl(area: float, speed: float) -> float:
    """计算断面积与风速对应的风量。"""
    return area * speed


def test_declaration_has_no_import_registration_and_alias_is_preserved():
    assert "airflow" not in tools.get_tool_names()
    catalog = ExtensionRegistry()
    catalog.register_agent(name="ventilation", description="处理通风计算", prompt="使用工具计算风量。", tools=[airflow_impl])
    catalog.validate()
    tools.install_specs(catalog.tools.values())
    tool = tools.get_all_tools()[0]
    assert tool.name == "airflow"
    assert tool.invoke({"area": 2, "speed": 3}) == 6
    from agents.prompts import build_system_prompt
    agents = catalog.subagents(set())
    assert agents[0]["system_prompt"] == "使用工具计算风量。"
    assert "处理通风计算" in build_system_prompt(agents)


@pytest.mark.asyncio
async def test_async_tool_alias():
    @define_tool(name="async_airflow")
    async def compute(area: float) -> float:
        """异步计算。"""
        return area * 2
    catalog = ExtensionRegistry()
    catalog.register_tool(compute)
    catalog.validate()
    tools.install_specs(catalog.tools.values())
    tool = tools.get_all_tools()[0]
    assert tool.name == "async_airflow"
    assert await tool.ainvoke({"area": 3}) == 6


def test_missing_reference_fails_and_forward_reference_works():
    catalog = ExtensionRegistry()
    catalog.register_agent(name="ventilation", description="通风", prompt="计算", tools=["airflow"])
    with pytest.raises(ValueError, match="未注册工具"):
        catalog.validate()
    catalog.register_tool(airflow_impl)
    catalog.validate()


@pytest.mark.parametrize("name", ["bad name", "../escape", "1abc", "", "a" * 65])
def test_invalid_names(name):
    with pytest.raises(ValueError, match="名称非法"):
        ExtensionRegistry().register_tool(airflow_impl, name=name)


def test_duplicate_and_reserved_names():
    catalog = ExtensionRegistry()
    catalog.register_tool(airflow_impl)
    with pytest.raises(ValueError, match="工具名冲突"):
        catalog.register_tool(airflow_impl)
    with pytest.raises(ValueError, match="内置工具冲突"):
        catalog.register_tool(airflow_impl, name="write_file")
    args = dict(name="ventilation", description="通风", prompt="计算", tools=["airflow"])
    catalog.register_agent(**args)
    with pytest.raises(ValueError, match="Agent 名称冲突"):
        catalog.register_agent(**args)
    with pytest.raises(ValueError, match="Agent 名称冲突"):
        catalog.register_agent(**{**args, "name": "general-purpose"})


@pytest.mark.parametrize("field", ["description", "prompt"])
def test_empty_prompts_and_descriptions(field):
    args = dict(name="ventilation", description="通风", prompt="计算", tools=[])
    args[field] = " "
    with pytest.raises(ValueError, match="不能为空"):
        ExtensionRegistry().register_agent(**args)


def test_approval_must_be_bound_and_configuration_must_exist():
    catalog = ExtensionRegistry()
    catalog.register_tool(airflow_impl)
    catalog.register_agent(name="ventilation", description="通风", prompt="计算", tools=[], approval_tools=["airflow"])
    with pytest.raises(ValueError, match="审批工具未绑定"):
        catalog.validate()
    with pytest.raises(ValueError, match="未注册工具"):
        ExtensionRegistry().validate(["typo"])


def test_shared_approval_applies_to_every_agent_using_tool():
    catalog = ExtensionRegistry()
    catalog.register_agent(name="one", description="一", prompt="一", tools=[airflow_impl], approval_tools=[airflow_impl])
    catalog.register_agent(name="two", description="二", prompt="二", tools=[airflow_impl])
    approvals = catalog.validate()
    tools.install_specs(catalog.tools.values())
    assert approvals == {"airflow"}
    assert all(agent["interrupt_on"] == {"airflow": True} for agent in catalog.subagents(approvals))


def test_failed_load_does_not_publish_partial_tools(monkeypatch):
    baseline = tools.get_tool_names()
    extensions.reset_for_tests()
    def setup(catalog):
        catalog.register_tool(airflow_impl)
        raise ValueError("broken extension")
    monkeypatch.setitem(sys.modules, "extensions.broken.extension", SimpleNamespace(setup=setup))
    monkeypatch.setattr(extensions, "discover_modules", lambda: ["extensions.broken.extension"])
    with pytest.raises(ValueError, match="extensions.broken.extension.*broken extension"):
        extensions.load_extensions()
    assert tools.get_tool_names() == baseline
    assert extensions._loaded is None


def test_missing_setup_fails_with_module_name(monkeypatch):
    extensions.reset_for_tests()
    monkeypatch.setitem(sys.modules, "extensions.broken.extension", SimpleNamespace())
    monkeypatch.setattr(extensions, "discover_modules", lambda: ["extensions.broken.extension"])
    with pytest.raises(ValueError, match="extensions.broken.extension.*setup"):
        extensions.load_extensions()


def test_multiple_cached_modules_reset_and_concurrent_load_are_idempotent(monkeypatch):
    base = extensions.discover_modules()
    calls = []
    def setup(catalog):
        calls.append(1)
        catalog.register_tool(airflow_impl)
    monkeypatch.setitem(sys.modules, "extensions.air.extension", SimpleNamespace(setup=setup))
    monkeypatch.setattr(extensions, "discover_modules", lambda: [*base, "extensions.air.extension"])
    for _ in range(2):
        tools.reset_for_tests()
        with ThreadPoolExecutor(max_workers=4) as pool:
            results = list(pool.map(lambda _: extensions.load_extensions(), range(8)))
        assert all(result is results[0] for result in results)
        assert len(tools.get_tool_names()) == 6
        assert "ask_user" in tools.get_tool_names() and "airflow" in tools.get_tool_names()
    assert len(calls) == 2
    with pytest.raises(ValueError, match="重启"):
        results[0].register_tool(airflow_impl)


def test_discovery_is_sorted_and_ignores_helpers(tmp_path, monkeypatch):
    monkeypatch.setattr(extensions, "__file__", str(tmp_path / "registry.py"))
    for name in ("zeta", "alpha", "_hidden"):
        folder = tmp_path / name
        folder.mkdir()
        (folder / "extension.py").write_text("def setup(registry): pass")
    (tmp_path / "helper.py").write_text("raise RuntimeError('must not import')")
    assert extensions.discover_modules() == ["extensions.alpha.extension", "extensions.zeta.extension"]
    (tmp_path / "linked").symlink_to(tmp_path / "alpha", target_is_directory=True)
    with pytest.raises(ValueError, match="符号链接"):
        extensions.discover_modules()


@pytest.mark.asyncio
async def test_invalid_extension_blocks_startup_before_database(monkeypatch):
    import main
    tools.reset_for_tests()
    monkeypatch.setenv("HITL_TOOLS", "unknown_tool")
    def forbidden():
        pytest.fail("扩展校验失败后不应初始化数据库")
    monkeypatch.setattr(main.db_conn, "init_schema", forbidden)
    with pytest.raises(ValueError, match="未注册工具"):
        await main._lifespan_startup()


def test_release_includes_extension_code_and_prompt_files():
    root = Path(__file__).resolve().parents[1]
    sys.path.insert(0, str(root / "release"))
    import build
    files = {p.relative_to(root).as_posix() for p in build.source_files(root)}
    assert "extensions/demo/extension.py" in files
    assert "extensions/core/extension.py" in files
    assert {".py", ".md", ".txt"} <= build.TREE_SUFFIXES["extensions"]