195 lines
4.9 KiB
Python
195 lines
4.9 KiB
Python
"""Tests for debate agent functionality."""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, MagicMock
|
|
from src.agent import DebateAgent
|
|
from src.constants import MAX_MEMORY_TOKENS
|
|
|
|
|
|
class MockProvider:
|
|
"""Mock LLM provider for testing."""
|
|
|
|
def __init__(self, model="test-model"):
|
|
self.model = model
|
|
|
|
def generate_response(self, messages, **kwargs):
|
|
return "This is a test response from the mock provider."
|
|
|
|
def validate_config(self):
|
|
return True
|
|
|
|
|
|
def test_agent_initialization():
|
|
"""Test basic agent initialization."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="You are a test agent"
|
|
)
|
|
|
|
assert agent.name == "Test Agent"
|
|
assert agent.provider == provider
|
|
assert agent.system_prompt == "You are a test agent"
|
|
assert len(agent.memory) == 0
|
|
|
|
|
|
def test_agent_set_position():
|
|
"""Test setting agent position."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt"
|
|
)
|
|
|
|
agent.set_position("for")
|
|
assert agent.position == "for"
|
|
|
|
agent.set_position("against")
|
|
assert agent.position == "against"
|
|
|
|
|
|
def test_agent_set_topic():
|
|
"""Test setting debate topic."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt"
|
|
)
|
|
|
|
agent.set_position("for")
|
|
agent.set_topic("AI is beneficial")
|
|
|
|
# Should create initial system message
|
|
assert len(agent.memory) == 1
|
|
assert agent.memory[0]["role"] == "system"
|
|
assert "AI is beneficial" in agent.memory[0]["content"]
|
|
assert "for" in agent.memory[0]["content"]
|
|
|
|
|
|
def test_agent_add_message():
|
|
"""Test adding messages to agent memory."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt",
|
|
max_memory_tokens=10000
|
|
)
|
|
|
|
agent.add_message("user", "This is a test message")
|
|
|
|
assert len(agent.memory) == 1
|
|
assert agent.memory[0]["role"] == "user"
|
|
assert agent.memory[0]["content"] == "This is a test message"
|
|
|
|
|
|
def test_agent_generate_response():
|
|
"""Test generating a response."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt"
|
|
)
|
|
|
|
agent.set_position("for")
|
|
agent.set_topic("Test topic")
|
|
|
|
agent.add_message("user", "Present your argument")
|
|
response = agent.generate_response()
|
|
|
|
# Response should be added to memory
|
|
assert "test response" in response.lower()
|
|
# Should have system message, user message, and assistant response
|
|
assert len(agent.memory) == 3
|
|
assert agent.memory[-1]["role"] == "assistant"
|
|
|
|
|
|
def test_agent_memory_stats():
|
|
"""Test getting memory statistics."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt",
|
|
max_memory_tokens=10000
|
|
)
|
|
|
|
agent.add_message("user", "Test message")
|
|
|
|
stats = agent.get_memory_stats()
|
|
|
|
assert "message_count" in stats
|
|
assert "current_tokens" in stats
|
|
assert "max_tokens" in stats
|
|
assert "percentage" in stats
|
|
assert stats["message_count"] == 1
|
|
assert stats["max_tokens"] == 10000
|
|
|
|
|
|
def test_agent_clear_memory():
|
|
"""Test clearing agent memory."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt"
|
|
)
|
|
|
|
agent.add_message("user", "Message 1")
|
|
agent.add_message("user", "Message 2")
|
|
|
|
assert len(agent.memory) == 2
|
|
|
|
agent.clear_memory()
|
|
|
|
assert len(agent.memory) == 0
|
|
|
|
|
|
def test_agent_memory_truncation():
|
|
"""Test that memory is truncated when exceeding limits."""
|
|
provider = MockProvider()
|
|
# Set very small memory limit for testing
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt",
|
|
max_memory_tokens=100 # Very small limit
|
|
)
|
|
|
|
agent.set_topic("Test topic")
|
|
|
|
# Add many messages to exceed limit
|
|
for i in range(20):
|
|
agent.add_message("user", f"This is test message number {i} " * 10)
|
|
|
|
# Memory should have been truncated
|
|
# System message should still be there
|
|
assert any(msg["role"] == "system" for msg in agent.memory)
|
|
|
|
# Total tokens should be under limit
|
|
stats = agent.get_memory_stats()
|
|
assert stats["current_tokens"] <= 100
|
|
|
|
|
|
def test_agent_repr():
|
|
"""Test agent string representation."""
|
|
provider = MockProvider()
|
|
agent = DebateAgent(
|
|
name="Test Agent",
|
|
provider=provider,
|
|
system_prompt="Test prompt"
|
|
)
|
|
|
|
agent.set_position("for")
|
|
agent.add_message("user", "Test")
|
|
|
|
repr_str = repr(agent)
|
|
|
|
assert "Test Agent" in repr_str
|
|
assert "for" in repr_str
|
|
assert "tokens" in repr_str
|