107 lines
3.2 KiB
Python
107 lines
3.2 KiB
Python
|
"""Tests for head registry plugin system."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from fusionagi.agents.head_registry import HeadRegistry, get_default_registry
|
||
|
|
|
||
|
|
|
||
|
|
class TestHeadRegistry:
|
||
|
|
def test_builtins_registered(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
assert reg.registered_count == 11 # 11 content heads (no Witness)
|
||
|
|
|
||
|
|
def test_create_builtin(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
head = reg.create("logic")
|
||
|
|
assert head._head_id.value == "logic"
|
||
|
|
|
||
|
|
def test_create_all(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
heads = reg.create_all()
|
||
|
|
assert len(heads) == 11
|
||
|
|
|
||
|
|
def test_create_missing_raises(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
with pytest.raises(KeyError, match="nonexistent"):
|
||
|
|
reg.create("nonexistent")
|
||
|
|
|
||
|
|
def test_list_heads(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
heads = reg.list_heads()
|
||
|
|
assert len(heads) == 11
|
||
|
|
assert all(h["builtin"] for h in heads)
|
||
|
|
|
||
|
|
def test_register_custom(self) -> None:
|
||
|
|
from fusionagi.agents.head_agent import HeadAgent
|
||
|
|
from fusionagi.schemas.head import HeadId
|
||
|
|
|
||
|
|
reg = HeadRegistry()
|
||
|
|
|
||
|
|
def my_factory(adapter=None, **kwargs):
|
||
|
|
return HeadAgent(
|
||
|
|
head_id=HeadId.LOGIC,
|
||
|
|
role="Custom",
|
||
|
|
objective="Custom analysis",
|
||
|
|
system_prompt="You are custom.",
|
||
|
|
)
|
||
|
|
|
||
|
|
reg.register(
|
||
|
|
"custom_head",
|
||
|
|
role="Custom",
|
||
|
|
objective="Custom analysis",
|
||
|
|
factory=my_factory,
|
||
|
|
tags=["custom"],
|
||
|
|
)
|
||
|
|
assert reg.registered_count == 12
|
||
|
|
head = reg.create("custom_head")
|
||
|
|
assert head.role == "Custom"
|
||
|
|
|
||
|
|
def test_register_factory_decorator(self) -> None:
|
||
|
|
from fusionagi.agents.head_agent import HeadAgent
|
||
|
|
from fusionagi.schemas.head import HeadId
|
||
|
|
|
||
|
|
reg = HeadRegistry()
|
||
|
|
|
||
|
|
@reg.register_factory("decorated_head", role="Decorated", objective="Test")
|
||
|
|
def make_head(adapter=None, **kwargs):
|
||
|
|
return HeadAgent(
|
||
|
|
head_id=HeadId.LOGIC,
|
||
|
|
role="Decorated",
|
||
|
|
objective="Test",
|
||
|
|
system_prompt="Test",
|
||
|
|
)
|
||
|
|
|
||
|
|
assert "decorated_head" in [h["head_id"] for h in reg.list_heads()]
|
||
|
|
|
||
|
|
def test_unregister(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
assert reg.unregister("logic")
|
||
|
|
assert reg.registered_count == 10
|
||
|
|
assert not reg.unregister("logic")
|
||
|
|
|
||
|
|
def test_create_all_with_tags(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
heads = reg.create_all(include_tags=["builtin"])
|
||
|
|
assert len(heads) == 11
|
||
|
|
heads_none = reg.create_all(include_tags=["nonexistent"])
|
||
|
|
assert len(heads_none) == 0
|
||
|
|
|
||
|
|
def test_get_spec(self) -> None:
|
||
|
|
reg = HeadRegistry()
|
||
|
|
spec = reg.get_spec("logic")
|
||
|
|
assert spec is not None
|
||
|
|
assert spec.role == "Logic"
|
||
|
|
assert reg.get_spec("nonexistent") is None
|
||
|
|
|
||
|
|
def test_no_auto_register(self) -> None:
|
||
|
|
reg = HeadRegistry(auto_register_builtins=False)
|
||
|
|
assert reg.registered_count == 0
|
||
|
|
|
||
|
|
|
||
|
|
class TestDefaultRegistry:
|
||
|
|
def test_get_default_registry(self) -> None:
|
||
|
|
reg = get_default_registry()
|
||
|
|
assert reg.registered_count >= 11
|