tests/test_agent_core.py

import pytest
from gaia_agent import ModularGAIAAgent
import os

@pytest.fixture
def agent():
    return ModularGAIAAgent()

def test_tool_registry(agent):
    assert 'llama3_chat' in agent.tools
    assert 'extractive_qa' in agent.tools
    assert 'youtube_video_qa' in agent.tools

def test_fetch_questions_api(monkeypatch, agent):
    class MockResponse:
        def json(self):
            return [{"task_id": "1", "question": "What is 2+2?", "file_name": ""}]
        def raise_for_status(self):
            pass
    monkeypatch.setattr("requests.get", lambda url: MockResponse())
    questions = agent.fetch_questions(from_api=True)
    assert isinstance(questions, list)
    assert questions[0]["question"] == "What is 2+2?"

def test_download_file(monkeypatch, agent, tmp_path):
    test_file = tmp_path / "test.txt"
    monkeypatch.setattr("requests.get", lambda url: type("R", (), {"status_code": 200, "content": b"hello"})())
    fname = agent.download_file("testid", str(test_file))
    assert os.path.exists(fname)
    with open(fname) as f:
        assert f.read() == "hello"

def test_end_to_end(monkeypatch, agent):
    # Mock API and tools for a simple run
    monkeypatch.setattr(agent, "fetch_questions", lambda from_api, questions_path=None: [{"task_id": "1", "question": "What is 2+2?", "file_name": ""}])
    agent.tools['llama3_chat'] = lambda prompt: "4"
    results = agent.run(from_api=True)
    assert results[0]["answer"] == "4"