tests/unit/agents/test_base.py

# 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 == []