agents/central_ai_hub_test.py

import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import pytest
import asyncio
from core.central_ai_hub import CentralAIHub
from utils.llm_orchestrator import LLMOrchestrator

@pytest.fixture
async def hub():
    """Fixture to create a CentralAIHub instance."""
    api_key = os.getenv('LLM_API_KEY')
    hub = CentralAIHub(api_key)
    await hub.start()
    yield hub
    await hub.shutdown()

@pytest.mark.asyncio
async def test_system_initialization(hub):
    """Test system initialization and LLM connection."""
    assert hub.llm_orchestrator is not None
    assert isinstance(hub.llm_orchestrator, LLMOrchestrator)
    assert len(hub.agents) > 0

@pytest.mark.asyncio
async def test_task_delegation(hub):
    """Test task delegation functionality."""
    task = {
        'type': 'code_analysis',
        'content': 'def hello(): print("Hello, World!")',
        'requirements': ['code understanding', 'static analysis']
    }
    task_id = await hub.delegate_task(task)
    assert task_id is not None
    status = await hub.get_task_status(task_id)
    assert status['status'] in ['active', 'completed']

@pytest.mark.asyncio
async def test_agent_selection(hub):
    """Test agent selection logic."""
    task = {
        'type': 'code_generation',
        'content': 'Create a simple REST API',
        'requirements': ['code generation', 'API design']
    }
    agent_id = await hub.select_agent(task)
    assert agent_id in hub.agents

@pytest.mark.asyncio
async def test_error_handling(hub):
    """Test error handling capabilities."""
    with pytest.raises(ValueError):
        await hub.initialize_agent('non_existent_agent')

    with pytest.raises(Exception):
        await hub.delegate_task(None)

if __name__ == '__main__':
    pytest.main(['-v', 'central_ai_hub_test.py'])