File size: 11,188 Bytes
14c8d0a |
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 |
"""
Tests for agent lazy loading service
"""
import asyncio
import pytest
from datetime import datetime, timedelta
from unittest.mock import Mock, patch, AsyncMock
from src.services.agent_lazy_loader import (
AgentLazyLoader,
AgentMetadata,
agent_lazy_loader
)
from src.agents.deodoro import BaseAgent
from src.core.exceptions import AgentExecutionError
class MockAgent(BaseAgent):
"""Mock agent for testing."""
def __init__(self, name: str = "MockAgent"):
super().__init__()
self.name = name
self.initialized = False
async def initialize(self):
self.initialized = True
async def process(self, *args, **kwargs):
return {"result": "mock"}
class TestAgentLazyLoader:
"""Test agent lazy loader functionality."""
@pytest.fixture
async def lazy_loader(self):
"""Create a lazy loader instance."""
loader = AgentLazyLoader(
unload_after_minutes=1,
max_loaded_agents=3
)
await loader.start()
yield loader
await loader.stop()
async def test_register_agent(self, lazy_loader):
"""Test agent registration."""
# Register a new agent
lazy_loader.register_agent(
name="TestAgent",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Test agent",
capabilities=["testing"],
priority=5,
preload=False
)
# Check registration
assert "TestAgent" in lazy_loader._registry
metadata = lazy_loader._registry["TestAgent"]
assert metadata.name == "TestAgent"
assert metadata.module_path == "tests.unit.test_agent_lazy_loader"
assert metadata.class_name == "MockAgent"
assert metadata.description == "Test agent"
assert metadata.capabilities == ["testing"]
assert metadata.priority == 5
assert metadata.preload is False
async def test_get_agent_class_lazy_load(self, lazy_loader):
"""Test lazy loading of agent class."""
# Register agent
lazy_loader.register_agent(
name="TestAgent",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Test agent",
capabilities=["testing"]
)
# Get agent class (should trigger lazy load)
agent_class = await lazy_loader.get_agent_class("TestAgent")
# Verify
assert agent_class == MockAgent
assert lazy_loader._stats["cache_misses"] == 1
assert lazy_loader._stats["total_loads"] == 1
# Get again (should use cache)
agent_class2 = await lazy_loader.get_agent_class("TestAgent")
assert agent_class2 == MockAgent
assert lazy_loader._stats["cache_hits"] == 1
async def test_create_agent_instance(self, lazy_loader):
"""Test creating agent instance."""
# Register agent
lazy_loader.register_agent(
name="TestAgent",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Test agent",
capabilities=["testing"]
)
# Create instance
agent = await lazy_loader.create_agent("TestAgent")
# Verify
assert isinstance(agent, MockAgent)
assert agent.initialized # Should be initialized
assert len(lazy_loader._instances) == 1
async def test_agent_not_found(self, lazy_loader):
"""Test error when agent not found."""
with pytest.raises(AgentExecutionError) as exc_info:
await lazy_loader.get_agent_class("NonExistentAgent")
assert "not registered" in str(exc_info.value)
async def test_preload_agents(self, lazy_loader):
"""Test preloading of high-priority agents."""
# Register agents with different priorities
lazy_loader.register_agent(
name="HighPriority",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="High priority agent",
capabilities=["high"],
priority=10,
preload=True
)
lazy_loader.register_agent(
name="LowPriority",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Low priority agent",
capabilities=["low"],
priority=1,
preload=False
)
# Preload
await lazy_loader._preload_agents()
# Check only high priority is loaded
assert lazy_loader._registry["HighPriority"].loaded_class is not None
assert lazy_loader._registry["LowPriority"].loaded_class is None
async def test_memory_pressure_unloading(self, lazy_loader):
"""Test unloading agents under memory pressure."""
lazy_loader.max_loaded_agents = 2
# Register and load 3 agents
for i in range(3):
lazy_loader.register_agent(
name=f"Agent{i}",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description=f"Agent {i}",
capabilities=[f"cap{i}"],
preload=False
)
# Load with delay to ensure different timestamps
await lazy_loader.get_agent_class(f"Agent{i}")
await asyncio.sleep(0.1)
# Should have unloaded oldest agent
loaded_count = sum(
1 for m in lazy_loader._registry.values()
if m.loaded_class
)
assert loaded_count <= 2
async def test_cleanup_unused_agents(self, lazy_loader):
"""Test cleanup of unused agents."""
# Register and load agent
lazy_loader.register_agent(
name="UnusedAgent",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Unused agent",
capabilities=["unused"],
preload=False
)
await lazy_loader.get_agent_class("UnusedAgent")
# Set last used to past
metadata = lazy_loader._registry["UnusedAgent"]
metadata.last_used = datetime.now() - timedelta(minutes=2)
# Run cleanup
await lazy_loader._cleanup_unused_agents()
# Should be unloaded
assert metadata.loaded_class is None
assert lazy_loader._stats["total_unloads"] == 1
async def test_get_available_agents(self, lazy_loader):
"""Test getting available agents list."""
# Register some agents
lazy_loader.register_agent(
name="Agent1",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Agent 1",
capabilities=["cap1"],
priority=10
)
lazy_loader.register_agent(
name="Agent2",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Agent 2",
capabilities=["cap2"],
priority=5
)
# Get available agents
agents = lazy_loader.get_available_agents()
# Check sorted by priority
assert len(agents) >= 2
assert agents[0]["priority"] >= agents[1]["priority"]
# Check fields
agent1 = next(a for a in agents if a["name"] == "Agent1")
assert agent1["description"] == "Agent 1"
assert agent1["capabilities"] == ["cap1"]
assert agent1["priority"] == 10
async def test_get_stats(self, lazy_loader):
"""Test getting loader statistics."""
# Perform some operations
lazy_loader.register_agent(
name="StatsAgent",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Stats agent",
capabilities=["stats"]
)
await lazy_loader.get_agent_class("StatsAgent")
await lazy_loader.create_agent("StatsAgent")
# Get stats
stats = lazy_loader.get_stats()
# Verify
assert stats["total_agents"] >= 1
assert stats["loaded_agents"] >= 1
assert stats["active_instances"] >= 1
assert stats["statistics"]["total_loads"] >= 1
assert stats["memory_usage"]["max_loaded_agents"] == 3
async def test_unload_with_active_instances(self, lazy_loader):
"""Test that agents with active instances are not unloaded."""
# Register and create agent
lazy_loader.register_agent(
name="ActiveAgent",
module_path="tests.unit.test_agent_lazy_loader",
class_name="MockAgent",
description="Active agent",
capabilities=["active"]
)
# Create instance (keeps reference)
agent = await lazy_loader.create_agent("ActiveAgent")
# Try to unload
metadata = lazy_loader._registry["ActiveAgent"]
await lazy_loader._unload_agent(metadata)
# Should still be loaded due to active instance
assert metadata.loaded_class is not None
async def test_invalid_module_path(self, lazy_loader):
"""Test loading agent with invalid module path."""
# Register with invalid module
lazy_loader.register_agent(
name="InvalidAgent",
module_path="invalid.module.path",
class_name="MockAgent",
description="Invalid agent",
capabilities=["invalid"]
)
# Should raise error
with pytest.raises(AgentExecutionError) as exc_info:
await lazy_loader.get_agent_class("InvalidAgent")
assert "Failed to load agent" in str(exc_info.value)
async def test_invalid_class_name(self, lazy_loader):
"""Test loading agent with invalid class name."""
# Register with invalid class
lazy_loader.register_agent(
name="InvalidClass",
module_path="tests.unit.test_agent_lazy_loader",
class_name="NonExistentClass",
description="Invalid class",
capabilities=["invalid"]
)
# Should raise error
with pytest.raises(AgentExecutionError):
await lazy_loader.get_agent_class("InvalidClass")
@pytest.mark.asyncio
async def test_global_lazy_loader():
"""Test global lazy loader instance."""
from src.services.agent_lazy_loader import get_agent_lazy_loader
# Get instance
loader = await get_agent_lazy_loader()
# Should be started
assert loader._running
# Should be same instance
loader2 = await get_agent_lazy_loader()
assert loader is loader2
# Cleanup
await loader.stop() |