python/tests/server/api_routes/test_agents_api.py

from unittest.mock import AsyncMock, patch

from fastapi.testclient import TestClient

from src.server.api_routes.agents_api import get_current_user
from src.server.services.agent_service import AI_AGENT_ROLES  # Import the actual roles

# It's important to patch the service BEFORE it's imported by the app
# so that the router gets the mocked dependency
mock_agent_service = AsyncMock()
mock_agent_service.get_assignable_agents.return_value = [
    {"id": agent_id, "name": role_name, "role": role_name} for role_name, agent_id in AI_AGENT_ROLES.items()
]

# The patch needs to target where the object is *used*, which is in the api_routes module
module_patch = patch("src.server.api_routes.agents_api.agent_service", mock_agent_service)

# Now we can import the app
from src.server.main import app  # noqa: E402

# Start the patch AFTER the app is imported
module_patch.start()

client = TestClient(app)


def test_get_assignable_agents_success():
    """
    Test successfully fetching the list of assignable agents.
    """
    # Mock authentication for admin
    app.dependency_overrides[get_current_user] = lambda: {"id": "admin-1", "role": "system_admin"}

    response = client.get("/api/agents/assignable")

    # Assertions
    assert response.status_code == 200

    data = response.json()
    assert isinstance(data, list)
    assert len(data) == len(AI_AGENT_ROLES)
    # Check that the returned agents match the expected roles
    expected_agents = [
        {"id": agent_id, "name": role_name, "role": role_name} for role_name, agent_id in AI_AGENT_ROLES.items()
    ]
    assert all(any(item == expected for expected in expected_agents) for item in data)

    # Verify that the mocked service method was called
    mock_agent_service.get_assignable_agents.assert_awaited_once()


def test_get_assignable_agents_service_error():
    """
    Test handling of an error from the agent service.
    """
    # Mock authentication
    app.dependency_overrides[get_current_user] = lambda: {"id": "admin-1", "role": "system_admin"}

    # Reset and configure the mock for this specific test
    mock_agent_service.reset_mock()
    mock_agent_service.get_assignable_agents.side_effect = Exception("Service unavailable")

    response = client.get("/api/agents/assignable")

    # Assert that the app's exception handler caught it (FastAPI returns 500)
    assert response.status_code == 500
    assert "Failed to retrieve assignable agents" in response.text

    # Clean up the side effect
    mock_agent_service.get_assignable_agents.side_effect = None
    mock_agent_service.reset_mock()
    # Restore the original return value for other tests if needed
    mock_agent_service.get_assignable_agents.return_value = [
        {"id": "ai-tester-1", "name": "測試 AI Agent", "role": "TEST"},
    ]


# Stop the patch after all tests in this module are done
def teardown_module(module):
    module_patch.stop()