test_app.py

from fastapi.testclient import TestClient
from app import app, main_agent

# TestClient is synchronous, so no need for pytest.mark.asyncio
client = TestClient(app)

def test_read_root():
    response = client.get("/")
    assert response.status_code == 200
    assert response.json() == {"message": "Hello, FastAPI!"}

def test_agent_chat_missing_query():
    response = client.post("/agent_chat", json={})
    assert response.status_code == 400
    assert response.json() == {"error": "Query parameter is missing"}

def test_agent_chat_with_query():
    # Mock the agent's run method to avoid actual API calls during testing
    original_run = main_agent.agent.run
    main_agent.agent.run = lambda x: "Mocked agent response for: " + x

    response = client.post("/agent_chat", json={"query": "Test query"})
    
    assert response.status_code == 200
    assert response.json() == {"response": "Mocked agent response for: Test query"}

    # Restore the original run method
    main_agent.agent.run = original_run