|
2 | 2 |
|
3 | 3 | from datetime import datetime, timezone |
4 | 4 | from unittest.mock import AsyncMock, patch |
| 5 | +from uuid import UUID |
| 6 | + |
| 7 | +from fastmcp import Client |
5 | 8 |
|
6 | 9 | from prefect_mcp_server._prefect_client.rate_limits import get_rate_limits |
| 10 | +from prefect_mcp_server.server import build_prefect_mcp_server |
7 | 11 |
|
8 | 12 |
|
9 | 13 | async def test_get_rate_limits_empty_response() -> None: |
@@ -816,3 +820,56 @@ async def test_get_rate_limits_handles_errors() -> None: |
816 | 820 | assert result["account_id"] is None |
817 | 821 | assert result["summary"] is None |
818 | 822 | assert result["throttling_periods"] is None |
| 823 | + |
| 824 | + |
| 825 | +async def test_get_rate_limits_passes_workspace_id_to_clients() -> None: |
| 826 | + """workspace_id must reach both clients so OAuth mode can scope the request.""" |
| 827 | + workspace_id = UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb") |
| 828 | + |
| 829 | + mock_client = AsyncMock() |
| 830 | + mock_client.api_url = ( |
| 831 | + f"https://api.prefect.cloud/api/accounts/abc-123/workspaces/{workspace_id}" |
| 832 | + ) |
| 833 | + |
| 834 | + mock_cloud_client = AsyncMock() |
| 835 | + mock_cloud_client.get = AsyncMock( |
| 836 | + return_value={ |
| 837 | + "account": "abc-123", |
| 838 | + "since": "2025-09-30T00:00:00Z", |
| 839 | + "until": "2025-10-01T00:00:00Z", |
| 840 | + "minutes": [], |
| 841 | + "keys": {}, |
| 842 | + } |
| 843 | + ) |
| 844 | + |
| 845 | + with ( |
| 846 | + patch( |
| 847 | + "prefect_mcp_server._prefect_client.rate_limits.get_prefect_client" |
| 848 | + ) as mock_get_client, |
| 849 | + patch( |
| 850 | + "prefect_mcp_server._prefect_client.rate_limits.get_prefect_cloud_client" |
| 851 | + ) as mock_get_cloud_client, |
| 852 | + ): |
| 853 | + mock_get_client.return_value.__aenter__.return_value = mock_client |
| 854 | + mock_get_cloud_client.return_value.__aenter__.return_value = mock_cloud_client |
| 855 | + mock_get_cloud_client.return_value.__aexit__.return_value = None |
| 856 | + |
| 857 | + result = await get_rate_limits(workspace_id=workspace_id) |
| 858 | + |
| 859 | + assert result["success"] is True |
| 860 | + mock_get_client.assert_called_once_with(workspace_id=workspace_id) |
| 861 | + mock_get_cloud_client.assert_called_once_with(workspace_id=workspace_id) |
| 862 | + |
| 863 | + |
| 864 | +async def test_review_rate_limits_tool_accepts_workspace_id() -> None: |
| 865 | + """The MCP tool schema must expose workspace_id for OAuth mode.""" |
| 866 | + server = build_prefect_mcp_server( |
| 867 | + include_docs_proxy=False, |
| 868 | + include_cloud_tools=True, |
| 869 | + ) |
| 870 | + |
| 871 | + async with Client(server) as client: |
| 872 | + tools = {tool.name: tool for tool in await client.list_tools()} |
| 873 | + |
| 874 | + schema = tools["review_rate_limits"].inputSchema |
| 875 | + assert "workspace_id" in schema["properties"] |
0 commit comments