feat: implement permanent tools (calculator, date/time, search)
Add three permanent tools for Tatlock agent: - Calculator: Safe math expression evaluation (arithmetic, algebra, trig, log) - Date/Time toolkit: Current time, relative dates, time differences - Web Search: SearXNG integration for privacy-preserving search Tools use PydanticAI @agent.tool decorator pattern with: - Clear docstrings visible to LLM - Error handling with string-based messages - Async support for I/O operations (web search) - 26 comprehensive tool tests
This commit is contained in:
@@ -0,0 +1,371 @@
|
||||
"""
|
||||
Tests for Tatlock's permanent tools (calculator, date/time, search).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from src.agents.tools import (
|
||||
calculate,
|
||||
get_current_datetime,
|
||||
calculate_time_offset,
|
||||
time_difference,
|
||||
search_web,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Calculator Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestCalculator:
|
||||
"""Tests for the calculator tool."""
|
||||
|
||||
def test_basic_arithmetic(self):
|
||||
"""Test basic arithmetic operations."""
|
||||
assert calculate("2 + 2") == "4"
|
||||
assert calculate("10 - 3") == "7"
|
||||
assert calculate("5 * 6") == "30"
|
||||
assert calculate("20 / 4") == "5" # Integer result, no decimal
|
||||
|
||||
def test_complex_expressions(self):
|
||||
"""Test complex mathematical expressions."""
|
||||
assert calculate("(2 + 3) * 4") == "20"
|
||||
assert calculate("10 ** 2") == "100"
|
||||
assert calculate("17 % 5") == "2"
|
||||
|
||||
def test_math_functions(self):
|
||||
"""Test mathematical functions."""
|
||||
assert calculate("sqrt(16)") == "4" # Integer result
|
||||
assert calculate("abs(-5)") == "5"
|
||||
assert calculate("round(3.7)") == "4"
|
||||
|
||||
# Test with constants
|
||||
result = calculate("pi * 2")
|
||||
assert "6.28" in result # Approximately 6.283...
|
||||
|
||||
def test_trigonometry(self):
|
||||
"""Test trigonometric functions."""
|
||||
result = calculate("sin(0)")
|
||||
assert result == "0" # Integer result
|
||||
|
||||
# cos(0) should be 1
|
||||
result = calculate("cos(0)")
|
||||
assert result == "1" # Integer result
|
||||
|
||||
def test_logarithms(self):
|
||||
"""Test logarithmic functions."""
|
||||
result = calculate("log10(100)")
|
||||
assert result == "2" # Integer result
|
||||
|
||||
result = calculate("exp(0)")
|
||||
assert result == "1" # Integer result
|
||||
|
||||
def test_error_handling(self):
|
||||
"""Test error handling for invalid expressions."""
|
||||
result = calculate("1 / 0")
|
||||
assert "Error: Division by zero" in result
|
||||
|
||||
result = calculate("invalid_function(5)")
|
||||
assert "Error calculating" in result
|
||||
|
||||
def test_integer_results(self):
|
||||
"""Test that integer results don't show unnecessary decimals."""
|
||||
assert calculate("4.0 + 6.0") == "10"
|
||||
assert calculate("sqrt(9)") == "3"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Date/Time Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestDateTime:
|
||||
"""Tests for date/time toolkit."""
|
||||
|
||||
def test_get_current_datetime_full(self):
|
||||
"""Test getting full current datetime."""
|
||||
result = get_current_datetime("full")
|
||||
# Should match format YYYY-MM-DD HH:MM:SS
|
||||
assert len(result) == 19
|
||||
assert result[4] == "-"
|
||||
assert result[7] == "-"
|
||||
assert result[10] == " "
|
||||
assert result[13] == ":"
|
||||
assert result[16] == ":"
|
||||
|
||||
def test_get_current_datetime_date(self):
|
||||
"""Test getting current date only."""
|
||||
result = get_current_datetime("date")
|
||||
# Should match format YYYY-MM-DD
|
||||
assert len(result) == 10
|
||||
assert result[4] == "-"
|
||||
assert result[7] == "-"
|
||||
|
||||
# Verify it's a valid date
|
||||
datetime.strptime(result, "%Y-%m-%d")
|
||||
|
||||
def test_get_current_datetime_time(self):
|
||||
"""Test getting current time only."""
|
||||
result = get_current_datetime("time")
|
||||
# Should match format HH:MM:SS
|
||||
assert len(result) == 8
|
||||
assert result[2] == ":"
|
||||
assert result[5] == ":"
|
||||
|
||||
def test_get_current_datetime_iso(self):
|
||||
"""Test getting ISO format."""
|
||||
result = get_current_datetime("iso")
|
||||
# Should be parseable as ISO format
|
||||
datetime.fromisoformat(result)
|
||||
|
||||
def test_calculate_time_offset_days(self):
|
||||
"""Test calculating time offsets in days."""
|
||||
result = calculate_time_offset("1 day ago")
|
||||
assert len(result) == 19 # YYYY-MM-DD HH:MM:SS
|
||||
|
||||
result = calculate_time_offset("2 days from now")
|
||||
assert len(result) == 19
|
||||
|
||||
def test_calculate_time_offset_weeks(self):
|
||||
"""Test calculating time offsets in weeks."""
|
||||
result = calculate_time_offset("1 week ago")
|
||||
assert len(result) == 19
|
||||
|
||||
result = calculate_time_offset("2 weeks from now")
|
||||
assert len(result) == 19
|
||||
|
||||
def test_calculate_time_offset_months(self):
|
||||
"""Test calculating time offsets in months."""
|
||||
result = calculate_time_offset("1 month ago")
|
||||
assert len(result) == 19
|
||||
|
||||
result = calculate_time_offset("3 months from now")
|
||||
assert len(result) == 19
|
||||
|
||||
def test_calculate_time_offset_years(self):
|
||||
"""Test calculating time offsets in years."""
|
||||
result = calculate_time_offset("1 year ago")
|
||||
assert len(result) == 19
|
||||
|
||||
result = calculate_time_offset("2 years from now")
|
||||
assert len(result) == 19
|
||||
|
||||
def test_calculate_time_offset_hours(self):
|
||||
"""Test calculating time offsets in hours."""
|
||||
result = calculate_time_offset("5 hours ago")
|
||||
assert len(result) == 19
|
||||
|
||||
result = calculate_time_offset("3 hours from now")
|
||||
assert len(result) == 19
|
||||
|
||||
def test_calculate_time_offset_invalid(self):
|
||||
"""Test error handling for invalid time offsets."""
|
||||
result = calculate_time_offset("invalid input")
|
||||
assert "Error" in result
|
||||
assert "Cannot parse" in result
|
||||
|
||||
def test_time_difference(self):
|
||||
"""Test calculating time difference."""
|
||||
result = time_difference("2024-01-01", "2024-01-15")
|
||||
assert "14 day" in result
|
||||
|
||||
def test_time_difference_with_now(self):
|
||||
"""Test time difference with 'now'."""
|
||||
# Get today's date
|
||||
today = datetime.now().strftime("%Y-%m-%d")
|
||||
result = time_difference(today, "now")
|
||||
# Should be less than a day
|
||||
assert "Less than" in result or "hour" in result or "minute" in result
|
||||
|
||||
def test_time_difference_with_times(self):
|
||||
"""Test time difference with full timestamps."""
|
||||
result = time_difference("2024-01-01 10:00:00", "2024-01-01 14:30:00")
|
||||
assert "4 hour" in result
|
||||
assert "30 minute" in result
|
||||
|
||||
def test_time_difference_error(self):
|
||||
"""Test error handling for invalid dates."""
|
||||
result = time_difference("invalid-date", "now")
|
||||
assert "Error" in result
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Search Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestSearch:
|
||||
"""Tests for web search tool."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_web_success(self):
|
||||
"""Test successful web search."""
|
||||
mock_response = {
|
||||
"results": [
|
||||
{
|
||||
"title": "Test Result 1",
|
||||
"url": "https://example.com/1",
|
||||
"content": "This is a test result"
|
||||
},
|
||||
{
|
||||
"title": "Test Result 2",
|
||||
"url": "https://example.com/2",
|
||||
"content": "Another test result"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
with patch("src.agents.tools.httpx.AsyncClient") as mock_client_class:
|
||||
# Create mock response
|
||||
mock_response_obj = type('MockResponse', (), {
|
||||
'status_code': 200,
|
||||
'json': lambda *args, **kwargs: mock_response
|
||||
})()
|
||||
|
||||
# Create mock client with async get method
|
||||
async def mock_get(*args, **kwargs):
|
||||
return mock_response_obj
|
||||
|
||||
mock_client_instance = type('MockClient', (), {
|
||||
'get': mock_get
|
||||
})()
|
||||
|
||||
# Setup async context manager
|
||||
async def mock_aenter(*args, **kwargs):
|
||||
return mock_client_instance
|
||||
|
||||
async def mock_aexit(*args, **kwargs):
|
||||
return None
|
||||
|
||||
mock_client_class.return_value.__aenter__ = mock_aenter
|
||||
mock_client_class.return_value.__aexit__ = mock_aexit
|
||||
|
||||
result = await search_web("test query", num_results=2)
|
||||
|
||||
assert "Test Result 1" in result
|
||||
assert "https://example.com/1" in result
|
||||
assert "Test Result 2" in result
|
||||
assert "https://example.com/2" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_web_no_results(self):
|
||||
"""Test web search with no results."""
|
||||
mock_response_data = {"results": []}
|
||||
|
||||
with patch("src.agents.tools.httpx.AsyncClient") as mock_client_class:
|
||||
mock_response_obj = type('MockResponse', (), {
|
||||
'status_code': 200,
|
||||
'json': lambda *args, **kwargs: mock_response_data
|
||||
})()
|
||||
|
||||
async def mock_get(*args, **kwargs):
|
||||
return mock_response_obj
|
||||
|
||||
mock_client_instance = type('MockClient', (), {
|
||||
'get': mock_get
|
||||
})()
|
||||
|
||||
async def mock_aenter(*args, **kwargs):
|
||||
return mock_client_instance
|
||||
|
||||
async def mock_aexit(*args, **kwargs):
|
||||
return None
|
||||
|
||||
mock_client_class.return_value.__aenter__ = mock_aenter
|
||||
mock_client_class.return_value.__aexit__ = mock_aexit
|
||||
|
||||
result = await search_web("test query")
|
||||
|
||||
assert "No results found" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_web_connection_error(self):
|
||||
"""Test web search with connection error."""
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_client_instance = AsyncMock()
|
||||
mock_client_instance.get.side_effect = Exception("Connection failed")
|
||||
mock_client.return_value.__aenter__.return_value = mock_client_instance
|
||||
|
||||
result = await search_web("test query")
|
||||
|
||||
assert "Error searching" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_web_limits_results(self):
|
||||
"""Test that search limits results to max 10."""
|
||||
mock_response_data = {
|
||||
"results": [
|
||||
{"title": f"Result {i}", "url": f"https://example.com/{i}", "content": "Test"}
|
||||
for i in range(20)
|
||||
]
|
||||
}
|
||||
|
||||
with patch("src.agents.tools.httpx.AsyncClient") as mock_client_class:
|
||||
mock_response_obj = type('MockResponse', (), {
|
||||
'status_code': 200,
|
||||
'json': lambda *args, **kwargs: mock_response_data
|
||||
})()
|
||||
|
||||
async def mock_get(*args, **kwargs):
|
||||
return mock_response_obj
|
||||
|
||||
mock_client_instance = type('MockClient', (), {
|
||||
'get': mock_get
|
||||
})()
|
||||
|
||||
async def mock_aenter(*args, **kwargs):
|
||||
return mock_client_instance
|
||||
|
||||
async def mock_aexit(*args, **kwargs):
|
||||
return None
|
||||
|
||||
mock_client_class.return_value.__aenter__ = mock_aenter
|
||||
mock_client_class.return_value.__aexit__ = mock_aexit
|
||||
|
||||
result = await search_web("test query", num_results=15)
|
||||
|
||||
# Should only return 10 results (max limit)
|
||||
result_count = result.count("URL:")
|
||||
assert result_count == 10
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_web_formats_results(self):
|
||||
"""Test that search results are properly formatted."""
|
||||
mock_response_data = {
|
||||
"results": [
|
||||
{
|
||||
"title": "Test Title",
|
||||
"url": "https://example.com",
|
||||
"content": "Test content description"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
with patch("src.agents.tools.httpx.AsyncClient") as mock_client_class:
|
||||
mock_response_obj = type('MockResponse', (), {
|
||||
'status_code': 200,
|
||||
'json': lambda *args, **kwargs: mock_response_data
|
||||
})()
|
||||
|
||||
async def mock_get(*args, **kwargs):
|
||||
return mock_response_obj
|
||||
|
||||
mock_client_instance = type('MockClient', (), {
|
||||
'get': mock_get
|
||||
})()
|
||||
|
||||
async def mock_aenter(*args, **kwargs):
|
||||
return mock_client_instance
|
||||
|
||||
async def mock_aexit(*args, **kwargs):
|
||||
return None
|
||||
|
||||
mock_client_class.return_value.__aenter__ = mock_aenter
|
||||
mock_client_class.return_value.__aexit__ = mock_aexit
|
||||
|
||||
result = await search_web("test query")
|
||||
|
||||
# Check formatting
|
||||
assert "1. Test Title" in result
|
||||
assert "URL: https://example.com" in result
|
||||
assert "Test content description" in result
|
||||
Reference in New Issue
Block a user