91e75e620b
CI: cua-driver distro-compat matrix / debian:12 (glibc 2.36) (push) Has been cancelled
CI: SPDX Headers / Check SPDX headers (warn-only) (push) Has been cancelled
CD: Docs MCP Server / build (linux/amd64) (push) Has been cancelled
CD: Docs MCP Server / build (linux/arm64) (push) Has been cancelled
CD: Docs MCP Server / merge (push) Has been cancelled
CI: cua-driver distro-compat matrix / Resolve release version (push) Has been cancelled
CI: cua-driver distro-compat matrix / fedora:41 (glibc 2.40) (push) Has been cancelled
CI: cua-driver distro-compat matrix / rockylinux:9 (glibc 2.34) (push) Has been cancelled
CI: cua-driver distro-compat matrix / ubuntu:22.04 (glibc 2.35) (push) Has been cancelled
CI: cua-driver distro-compat matrix / ubuntu:24.04 (glibc 2.39) (push) Has been cancelled
CI: cua-driver distro-compat matrix / Distro compat summary (push) Has been cancelled
CI: Rust Linux unit / Rust Linux unit and compile (push) Has been cancelled
CI: Rust Windows unit / Rust Windows unit and compile (push) Has been cancelled
CI: Nix Linux Rust source / Nix / compositor build (push) Has been cancelled
CI: Nix Linux Rust source / Nix / driver package (push) Has been cancelled
CI: Nix Linux Rust source / Nix / Rust unit tests (push) Has been cancelled
263 lines
9.1 KiB
Python
263 lines
9.1 KiB
Python
"""
|
|
Tests for centralized tool resolution in ComputerAgent.
|
|
|
|
Tests that:
|
|
1. FARA (specialized model) auto-wraps Computer to BrowserTool with warning
|
|
2. FARA accepts explicit BrowserTool without warning
|
|
3. Claude (general model) accepts any tool without wrapping
|
|
4. Custom function tools pass through unchanged
|
|
"""
|
|
|
|
import warnings
|
|
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
from cua_agent.computers import is_agent_computer
|
|
from cua_agent.tools.browser_tool import BrowserTool
|
|
from cua_agent.types import AgentConfigInfo
|
|
|
|
|
|
# Mock agent config class for testing
|
|
class MockAgentConfig:
|
|
"""Mock agent config class for testing"""
|
|
|
|
async def predict_step(self, **kwargs):
|
|
return {"output": [], "usage": {}}
|
|
|
|
async def predict_click(self, **kwargs):
|
|
return None
|
|
|
|
def get_capabilities(self):
|
|
return ["step"]
|
|
|
|
|
|
class TestToolResolution:
|
|
"""Tests for ComputerAgent._resolve_tools()"""
|
|
|
|
@pytest.fixture
|
|
def mock_computer(self):
|
|
"""Create a mock Computer object"""
|
|
computer = Mock()
|
|
computer.interface = Mock()
|
|
computer.interface.interface = Mock() # For hotkey, etc.
|
|
computer.interface.playwright_exec = Mock()
|
|
return computer
|
|
|
|
@pytest.fixture
|
|
def mock_browser_tool(self):
|
|
"""Create a mock BrowserTool"""
|
|
tool = Mock(spec=BrowserTool)
|
|
return tool
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fara_auto_wraps_computer_to_browser_tool(self, mock_computer):
|
|
"""FARA auto-wraps Computer to BrowserTool with warning."""
|
|
from cua_agent.agent import ComputerAgent
|
|
|
|
# Patch find_agent_config to return FARA config
|
|
fara_config = AgentConfigInfo(
|
|
agent_class=MockAgentConfig,
|
|
models_regex=r"(?i).*fara.*",
|
|
priority=0,
|
|
tool_type="browser",
|
|
)
|
|
|
|
# Patch is_agent_computer to recognize our mock as a computer
|
|
def mock_is_agent_computer(tool):
|
|
return tool is mock_computer
|
|
|
|
with patch("cua_agent.agent.find_agent_config", return_value=fara_config):
|
|
with patch("cua_agent.agent.is_agent_computer", side_effect=mock_is_agent_computer):
|
|
agent = ComputerAgent(model="cua/microsoft/fara-7b", tools=[mock_computer])
|
|
# Tool resolution happens in _resolve_tools, which is async
|
|
with pytest.warns(UserWarning, match="Auto-wrapping Computer to BrowserTool"):
|
|
resolved = await agent._resolve_tools(agent.tools, "browser")
|
|
|
|
# Should have wrapped to BrowserTool
|
|
assert len(resolved) == 1
|
|
assert isinstance(resolved[0], BrowserTool)
|
|
|
|
def test_fara_accepts_explicit_browser_tool_no_warning(self, mock_browser_tool):
|
|
"""FARA accepts explicit BrowserTool without warning."""
|
|
from cua_agent.agent import ComputerAgent
|
|
|
|
# Patch find_agent_config to return FARA config
|
|
fara_config = AgentConfigInfo(
|
|
agent_class=MockAgentConfig,
|
|
models_regex=r"(?i).*fara.*",
|
|
priority=0,
|
|
tool_type="browser",
|
|
)
|
|
|
|
with patch("cua_agent.agent.find_agent_config", return_value=fara_config):
|
|
# Should not raise any warnings
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
agent = ComputerAgent(model="cua/microsoft/fara-7b", tools=[mock_browser_tool])
|
|
|
|
# Should keep the original BrowserTool
|
|
assert len(agent.tools) == 1
|
|
assert agent.tools[0] is mock_browser_tool
|
|
|
|
def test_claude_accepts_any_tool_no_wrapping(self, mock_computer):
|
|
"""Claude (general model) accepts any tool without wrapping."""
|
|
from cua_agent.agent import ComputerAgent
|
|
|
|
# Patch find_agent_config to return Claude config (no tool_type)
|
|
claude_config = AgentConfigInfo(
|
|
agent_class=MockAgentConfig,
|
|
models_regex=r".*claude.*",
|
|
priority=0,
|
|
tool_type=None, # General model, no tool type requirement
|
|
)
|
|
|
|
with patch("cua_agent.agent.find_agent_config", return_value=claude_config):
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error") # Fail if any warning
|
|
agent = ComputerAgent(model="claude-sonnet-4", tools=[mock_computer])
|
|
|
|
# Should keep original Computer, not wrapped
|
|
assert len(agent.tools) == 1
|
|
assert agent.tools[0] is mock_computer
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_tools_pass_through(self, mock_computer):
|
|
"""Custom function tools pass through unchanged."""
|
|
from cua_agent.agent import ComputerAgent
|
|
|
|
def custom_tool():
|
|
"""A custom tool function."""
|
|
pass
|
|
|
|
# Patch find_agent_config to return FARA config
|
|
fara_config = AgentConfigInfo(
|
|
agent_class=MockAgentConfig,
|
|
models_regex=r"(?i).*fara.*",
|
|
priority=0,
|
|
tool_type="browser",
|
|
)
|
|
|
|
# Patch is_agent_computer to recognize our mock as a computer
|
|
def mock_is_agent_computer(tool):
|
|
return tool is mock_computer
|
|
|
|
with patch("cua_agent.agent.find_agent_config", return_value=fara_config):
|
|
with patch("cua_agent.agent.is_agent_computer", side_effect=mock_is_agent_computer):
|
|
agent = ComputerAgent(
|
|
model="cua/microsoft/fara-7b", tools=[mock_computer, custom_tool]
|
|
)
|
|
# Tool resolution happens in _resolve_tools, which is async
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
resolved = await agent._resolve_tools(agent.tools, "browser")
|
|
|
|
# Should have wrapped computer but kept custom tool unchanged
|
|
assert len(resolved) == 2
|
|
assert isinstance(resolved[0], BrowserTool)
|
|
assert resolved[1] is custom_tool
|
|
|
|
|
|
class TestAgentConfigInfo:
|
|
"""Tests for AgentConfigInfo with tool_type field"""
|
|
|
|
def test_agent_config_info_with_tool_type(self):
|
|
"""AgentConfigInfo accepts tool_type parameter"""
|
|
config = AgentConfigInfo(
|
|
agent_class=MockAgentConfig,
|
|
models_regex=r".*fara.*",
|
|
priority=0,
|
|
tool_type="browser",
|
|
)
|
|
assert config.tool_type == "browser"
|
|
|
|
def test_agent_config_info_without_tool_type(self):
|
|
"""AgentConfigInfo defaults tool_type to None"""
|
|
config = AgentConfigInfo(
|
|
agent_class=MockAgentConfig,
|
|
models_regex=r".*claude.*",
|
|
priority=0,
|
|
)
|
|
assert config.tool_type is None
|
|
|
|
def test_agent_config_info_matches_model(self):
|
|
"""AgentConfigInfo.matches_model works correctly"""
|
|
config = AgentConfigInfo(
|
|
agent_class=MockAgentConfig,
|
|
models_regex=r"(?i).*fara-7b.*",
|
|
priority=0,
|
|
tool_type="browser",
|
|
)
|
|
assert config.matches_model("cua/microsoft/fara-7b")
|
|
assert config.matches_model("FARA-7B")
|
|
assert not config.matches_model("claude-sonnet-4")
|
|
|
|
|
|
class TestRegisterAgentDecorator:
|
|
"""Tests for @register_agent decorator with tool_type"""
|
|
|
|
def test_register_agent_with_tool_type(self):
|
|
"""@register_agent accepts tool_type parameter"""
|
|
from cua_agent.decorators import _agent_configs, register_agent
|
|
|
|
# Clear registry for test
|
|
original_configs = _agent_configs.copy()
|
|
_agent_configs.clear()
|
|
|
|
try:
|
|
|
|
@register_agent(models=r"test-model.*", tool_type="browser")
|
|
class TestAgentConfig:
|
|
async def predict_step(self, **kwargs):
|
|
pass
|
|
|
|
async def predict_click(self, **kwargs):
|
|
pass
|
|
|
|
def get_capabilities(self):
|
|
return ["step"]
|
|
|
|
# Find the registered config
|
|
from cua_agent.decorators import find_agent_config
|
|
|
|
config = find_agent_config("test-model-123")
|
|
assert config is not None
|
|
assert config.tool_type == "browser"
|
|
|
|
finally:
|
|
# Restore original registry
|
|
_agent_configs.clear()
|
|
_agent_configs.extend(original_configs)
|
|
|
|
def test_register_agent_without_tool_type(self):
|
|
"""@register_agent without tool_type defaults to None"""
|
|
from cua_agent.decorators import _agent_configs, register_agent
|
|
|
|
# Clear registry for test
|
|
original_configs = _agent_configs.copy()
|
|
_agent_configs.clear()
|
|
|
|
try:
|
|
|
|
@register_agent(models=r"general-model.*")
|
|
class GeneralAgentConfig:
|
|
async def predict_step(self, **kwargs):
|
|
pass
|
|
|
|
async def predict_click(self, **kwargs):
|
|
pass
|
|
|
|
def get_capabilities(self):
|
|
return ["step"]
|
|
|
|
# Find the registered config
|
|
from cua_agent.decorators import find_agent_config
|
|
|
|
config = find_agent_config("general-model-123")
|
|
assert config is not None
|
|
assert config.tool_type is None
|
|
|
|
finally:
|
|
# Restore original registry
|
|
_agent_configs.clear()
|
|
_agent_configs.extend(original_configs)
|