69 lines
1.7 KiB
Python
69 lines
1.7 KiB
Python
"""Test the base class logic, message passing and background task handling."""
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
from control_backend.core.agent_system import AgentDirectory, BaseAgent, InternalMessage
|
|
|
|
|
|
class ConcreteTestAgent(BaseAgent):
|
|
logger = logging.getLogger("test")
|
|
|
|
def __init__(self, name: str):
|
|
super().__init__(name)
|
|
self.received = []
|
|
|
|
async def setup(self):
|
|
pass
|
|
|
|
async def handle_message(self, msg: InternalMessage):
|
|
self.received.append(msg)
|
|
if msg.body == "stop":
|
|
await self.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_lifecycle():
|
|
agent = ConcreteTestAgent("lifecycle_agent")
|
|
await agent.start()
|
|
assert agent._running is True
|
|
|
|
# Test background task
|
|
async def dummy_task():
|
|
await asyncio.sleep(0.01)
|
|
|
|
await agent.add_background_task(dummy_task())
|
|
assert len(agent._tasks) > 0
|
|
|
|
# Wait for task to finish
|
|
await asyncio.sleep(0.02)
|
|
assert len(agent._tasks) == 1 # _process_inbox is still running
|
|
|
|
await agent.stop()
|
|
assert agent._running is False
|
|
|
|
await asyncio.sleep(0.01)
|
|
|
|
# Tasks should be cancelled
|
|
assert len(agent._tasks) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_unknown_agent(caplog):
|
|
agent = ConcreteTestAgent("sender")
|
|
msg = InternalMessage(to="unknown_sender", sender="sender", body="boo")
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
await agent.send(msg)
|
|
|
|
assert "Attempted to send message to unknown agent: unknown_sender" in caplog.text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_agent():
|
|
agent = ConcreteTestAgent("registrant")
|
|
assert AgentDirectory.get("registrant") == agent
|
|
assert AgentDirectory.get("non_existent") is None
|