File size: 12,167 Bytes
56fd459 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 |
# Tests for ankigen_core/agents/base.py
import pytest
import asyncio
from unittest.mock import MagicMock, AsyncMock, patch
from dataclasses import dataclass
from typing import Dict, Any
from ankigen_core.agents.base import AgentConfig, BaseAgentWrapper, AgentResponse
# Test AgentConfig
def test_agent_config_creation():
"""Test basic AgentConfig creation"""
config = AgentConfig(
name="test_agent",
instructions="Test instructions",
model="gpt-4o",
temperature=0.7
)
assert config.name == "test_agent"
assert config.instructions == "Test instructions"
assert config.model == "gpt-4o"
assert config.temperature == 0.7
assert config.custom_prompts == {}
def test_agent_config_defaults():
"""Test AgentConfig with default values"""
config = AgentConfig(
name="test_agent",
instructions="Test instructions"
)
assert config.model == "gpt-4o"
assert config.temperature == 0.7
assert config.max_tokens is None
assert config.timeout == 30.0
assert config.retry_attempts == 3
assert config.enable_tracing is True
assert config.custom_prompts == {}
def test_agent_config_custom_prompts():
"""Test AgentConfig with custom prompts"""
custom_prompts = {"greeting": "Hello there", "farewell": "Goodbye"}
config = AgentConfig(
name="test_agent",
instructions="Test instructions",
custom_prompts=custom_prompts
)
assert config.custom_prompts == custom_prompts
# Test BaseAgentWrapper
@pytest.fixture
def mock_openai_client():
"""Mock OpenAI client for testing"""
return MagicMock()
@pytest.fixture
def test_agent_config():
"""Sample agent config for testing"""
return AgentConfig(
name="test_agent",
instructions="Test instructions",
model="gpt-4o",
temperature=0.7,
timeout=10.0,
retry_attempts=2
)
@pytest.fixture
def base_agent_wrapper(test_agent_config, mock_openai_client):
"""Base agent wrapper for testing"""
return BaseAgentWrapper(test_agent_config, mock_openai_client)
def test_base_agent_wrapper_init(base_agent_wrapper, test_agent_config, mock_openai_client):
"""Test BaseAgentWrapper initialization"""
assert base_agent_wrapper.config == test_agent_config
assert base_agent_wrapper.openai_client == mock_openai_client
assert base_agent_wrapper.agent is None
assert base_agent_wrapper.runner is None
assert base_agent_wrapper._performance_metrics == {
"total_calls": 0,
"successful_calls": 0,
"average_response_time": 0.0,
"error_count": 0,
}
@patch('ankigen_core.agents.base.Agent')
@patch('ankigen_core.agents.base.Runner')
async def test_base_agent_wrapper_initialize(mock_runner, mock_agent, base_agent_wrapper):
"""Test agent initialization"""
mock_agent_instance = MagicMock()
mock_runner_instance = MagicMock()
mock_agent.return_value = mock_agent_instance
mock_runner.return_value = mock_runner_instance
await base_agent_wrapper.initialize()
mock_agent.assert_called_once_with(
name="test_agent",
instructions="Test instructions",
model="gpt-4o",
temperature=0.7
)
mock_runner.assert_called_once_with(
agent=mock_agent_instance,
client=base_agent_wrapper.openai_client
)
assert base_agent_wrapper.agent == mock_agent_instance
assert base_agent_wrapper.runner == mock_runner_instance
@patch('ankigen_core.agents.base.Agent')
@patch('ankigen_core.agents.base.Runner')
async def test_base_agent_wrapper_initialize_error(mock_runner, mock_agent, base_agent_wrapper):
"""Test agent initialization with error"""
mock_agent.side_effect = Exception("Agent creation failed")
with pytest.raises(Exception, match="Agent creation failed"):
await base_agent_wrapper.initialize()
assert base_agent_wrapper.agent is None
assert base_agent_wrapper.runner is None
async def test_base_agent_wrapper_execute_without_initialization(base_agent_wrapper):
"""Test execute method when agent isn't initialized"""
with patch.object(base_agent_wrapper, 'initialize') as mock_init:
with patch.object(base_agent_wrapper, '_run_agent') as mock_run:
mock_run.return_value = "test response"
result = await base_agent_wrapper.execute("test input")
mock_init.assert_called_once()
mock_run.assert_called_once_with("test input")
assert result == "test response"
async def test_base_agent_wrapper_execute_with_context(base_agent_wrapper):
"""Test execute method with context"""
base_agent_wrapper.runner = MagicMock()
with patch.object(base_agent_wrapper, '_run_agent') as mock_run:
mock_run.return_value = "test response"
context = {"key1": "value1", "key2": "value2"}
result = await base_agent_wrapper.execute("test input", context)
expected_input = "test input\n\nContext:\nkey1: value1\nkey2: value2"
mock_run.assert_called_once_with(expected_input)
assert result == "test response"
async def test_base_agent_wrapper_execute_timeout(base_agent_wrapper):
"""Test execute method with timeout"""
base_agent_wrapper.runner = MagicMock()
with patch.object(base_agent_wrapper, '_run_agent') as mock_run:
mock_run.side_effect = asyncio.TimeoutError()
with pytest.raises(asyncio.TimeoutError):
await base_agent_wrapper.execute("test input")
assert base_agent_wrapper._performance_metrics["error_count"] == 1
async def test_base_agent_wrapper_execute_exception(base_agent_wrapper):
"""Test execute method with exception"""
base_agent_wrapper.runner = MagicMock()
with patch.object(base_agent_wrapper, '_run_agent') as mock_run:
mock_run.side_effect = Exception("Execution failed")
with pytest.raises(Exception, match="Execution failed"):
await base_agent_wrapper.execute("test input")
assert base_agent_wrapper._performance_metrics["error_count"] == 1
async def test_base_agent_wrapper_run_agent_success(base_agent_wrapper):
"""Test _run_agent method with successful execution"""
mock_runner = MagicMock()
mock_run = MagicMock()
mock_run.id = "run_123"
mock_run.status = "completed"
mock_run.thread_id = "thread_456"
mock_message = MagicMock()
mock_message.role = "assistant"
mock_message.content = "test response"
mock_runner.create_run = AsyncMock(return_value=mock_run)
mock_runner.get_run = AsyncMock(return_value=mock_run)
mock_runner.get_messages = AsyncMock(return_value=[mock_message])
base_agent_wrapper.runner = mock_runner
result = await base_agent_wrapper._run_agent("test input")
mock_runner.create_run.assert_called_once_with(
messages=[{"role": "user", "content": "test input"}]
)
mock_runner.get_messages.assert_called_once_with("thread_456")
assert result == "test response"
async def test_base_agent_wrapper_run_agent_retry(base_agent_wrapper):
"""Test _run_agent method with retry logic"""
mock_runner = MagicMock()
mock_runner.create_run = AsyncMock(side_effect=[
Exception("First attempt failed"),
Exception("Second attempt failed")
])
base_agent_wrapper.runner = mock_runner
with pytest.raises(Exception, match="Second attempt failed"):
await base_agent_wrapper._run_agent("test input")
assert mock_runner.create_run.call_count == 2
async def test_base_agent_wrapper_run_agent_no_response(base_agent_wrapper):
"""Test _run_agent method when no assistant response is found"""
mock_runner = MagicMock()
mock_run = MagicMock()
mock_run.id = "run_123"
mock_run.status = "completed"
mock_run.thread_id = "thread_456"
mock_message = MagicMock()
mock_message.role = "user" # No assistant response
mock_message.content = "user message"
mock_runner.create_run = AsyncMock(return_value=mock_run)
mock_runner.get_run = AsyncMock(return_value=mock_run)
mock_runner.get_messages = AsyncMock(return_value=[mock_message])
base_agent_wrapper.runner = mock_runner
with pytest.raises(ValueError, match="No assistant response found"):
await base_agent_wrapper._run_agent("test input")
def test_base_agent_wrapper_update_performance_metrics(base_agent_wrapper):
"""Test performance metrics update"""
base_agent_wrapper._update_performance_metrics(1.5, success=True)
metrics = base_agent_wrapper._performance_metrics
assert metrics["successful_calls"] == 1
assert metrics["average_response_time"] == 1.5
# Add another successful call
base_agent_wrapper._update_performance_metrics(2.5, success=True)
metrics = base_agent_wrapper._performance_metrics
assert metrics["successful_calls"] == 2
assert metrics["average_response_time"] == 2.0 # (1.5 + 2.5) / 2
def test_base_agent_wrapper_get_performance_metrics(base_agent_wrapper):
"""Test getting performance metrics"""
base_agent_wrapper._performance_metrics = {
"total_calls": 10,
"successful_calls": 8,
"average_response_time": 1.2,
"error_count": 2,
}
metrics = base_agent_wrapper.get_performance_metrics()
assert metrics["total_calls"] == 10
assert metrics["successful_calls"] == 8
assert metrics["average_response_time"] == 1.2
assert metrics["error_count"] == 2
assert metrics["success_rate"] == 0.8
assert metrics["agent_name"] == "test_agent"
async def test_base_agent_wrapper_handoff_to(base_agent_wrapper):
"""Test handoff to another agent"""
target_agent = MagicMock()
target_agent.config.name = "target_agent"
target_agent.execute = AsyncMock(return_value="handoff result")
context = {
"reason": "Test handoff",
"user_input": "Continue with this",
"additional_data": "some data"
}
result = await base_agent_wrapper.handoff_to(target_agent, context)
expected_context = {
"from_agent": "test_agent",
"handoff_reason": "Test handoff",
"user_input": "Continue with this",
"additional_data": "some data"
}
target_agent.execute.assert_called_once_with("Continue with this", expected_context)
assert result == "handoff result"
async def test_base_agent_wrapper_handoff_to_default_input(base_agent_wrapper):
"""Test handoff to another agent with default input"""
target_agent = MagicMock()
target_agent.config.name = "target_agent"
target_agent.execute = AsyncMock(return_value="handoff result")
context = {"reason": "Test handoff"}
result = await base_agent_wrapper.handoff_to(target_agent, context)
expected_context = {
"from_agent": "test_agent",
"handoff_reason": "Test handoff",
"reason": "Test handoff"
}
target_agent.execute.assert_called_once_with("Continue processing", expected_context)
assert result == "handoff result"
# Test AgentResponse
def test_agent_response_creation():
"""Test AgentResponse creation"""
response = AgentResponse(
success=True,
data={"cards": []},
agent_name="test_agent",
execution_time=1.5,
metadata={"version": "1.0"},
errors=["minor warning"]
)
assert response.success is True
assert response.data == {"cards": []}
assert response.agent_name == "test_agent"
assert response.execution_time == 1.5
assert response.metadata == {"version": "1.0"}
assert response.errors == ["minor warning"]
def test_agent_response_defaults():
"""Test AgentResponse with default values"""
response = AgentResponse(
success=True,
data={"result": "success"},
agent_name="test_agent",
execution_time=1.0
)
assert response.metadata == {}
assert response.errors == [] |