Spaces:
Runtime error
Runtime error
File size: 3,004 Bytes
0b2427a 87f4f6f 0b2427a a745005 0b2427a a745005 0b2427a a745005 0b2427a a745005 0b2427a a745005 0b2427a | 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 | """Unit tests for base agent class."""
import pytest
from unittest.mock import MagicMock, patch
from src.agents.base import BaseAgent
from src.utils.cost_tracker import CostTracker
class MockAgent(BaseAgent):
"""Concrete test agent for testing base class."""
def get_system_prompt(self) -> str:
return "Test system prompt"
async def run(self, **kwargs):
return {"result": "test"}
@pytest.mark.asyncio
async def test_base_agent_initialization():
"""Test base agent initializes correctly."""
tracker = CostTracker()
with patch("src.agents.base.get_settings") as mock_settings:
mock_settings.return_value = MagicMock(
default_model="x-ai/grok-4.1-fast:free",
openrouter_api_key="test-key",
openrouter_base_url="https://openrouter.ai/api/v1",
)
agent = MockAgent(
name="TestAgent",
model="openai/gpt-5-mini",
temperature=0.5,
cost_tracker=tracker,
)
assert agent.name == "TestAgent"
assert agent.model_name == "openai/gpt-5-mini"
assert agent.cost_tracker == tracker
@pytest.mark.asyncio
async def test_base_agent_uses_default_model():
"""Test agent uses default model from config."""
with patch("src.agents.base.get_settings") as mock_settings:
mock_settings.return_value = MagicMock(
default_model="x-ai/grok-4.1-fast:free",
openrouter_api_key="test-key",
openrouter_base_url="https://openrouter.ai/api/v1",
)
agent = MockAgent(name="TestAgent")
assert agent.model_name == "x-ai/grok-4.1-fast:free"
@pytest.mark.asyncio
async def test_create_messages():
"""Test message creation."""
with patch("src.agents.base.get_settings") as mock_settings:
mock_settings.return_value = MagicMock(
default_model="test-model",
openrouter_api_key="test-key",
openrouter_base_url="https://test.com",
)
agent = MockAgent(name="TestAgent")
messages = agent._create_messages("test user message")
assert len(messages) == 2
assert messages[0].content == "Test system prompt"
assert messages[1].content == "test user message"
@pytest.mark.asyncio
async def test_get_cost_summary():
"""Test cost summary retrieval."""
tracker = CostTracker()
with patch("src.agents.base.get_settings") as mock_settings:
mock_settings.return_value = MagicMock(
default_model="test-model",
openrouter_api_key="test-key",
openrouter_base_url="https://test.com",
)
agent = MockAgent(name="TestAgent", cost_tracker=tracker)
# Track some usage
tracker.track_usage("openai/gpt-5-mini", 1000, 500)
summary = agent.get_cost_summary()
assert summary["total_input_tokens"] == 1000
assert summary["total_output_tokens"] == 500
assert summary["calls"] == 1
|