Marcelo Trylesinski
122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
"""Tests for poll_task async iterator."""
|
|
|
|
from collections.abc import Callable, Coroutine
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from mcp.client.experimental.tasks import ExperimentalClientFeatures
|
|
from mcp.types import GetTaskResult, TaskStatus
|
|
|
|
|
|
def make_task_result(
|
|
status: TaskStatus = "working",
|
|
poll_interval: int = 0,
|
|
task_id: str = "test-task",
|
|
status_message: str | None = None,
|
|
) -> GetTaskResult:
|
|
"""Create GetTaskResult with sensible defaults."""
|
|
now = datetime.now(timezone.utc)
|
|
return GetTaskResult(
|
|
task_id=task_id,
|
|
status=status,
|
|
status_message=status_message,
|
|
created_at=now,
|
|
last_updated_at=now,
|
|
ttl=60000,
|
|
poll_interval=poll_interval,
|
|
)
|
|
|
|
|
|
def make_status_sequence(
|
|
*statuses: TaskStatus,
|
|
task_id: str = "test-task",
|
|
) -> Callable[[str], Coroutine[Any, Any, GetTaskResult]]:
|
|
"""Create mock get_task that returns statuses in sequence."""
|
|
status_iter = iter(statuses)
|
|
|
|
async def mock_get_task(tid: str) -> GetTaskResult:
|
|
return make_task_result(status=next(status_iter), task_id=tid)
|
|
|
|
return mock_get_task
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_session() -> AsyncMock:
|
|
return AsyncMock()
|
|
|
|
|
|
@pytest.fixture
|
|
def features(mock_session: AsyncMock) -> ExperimentalClientFeatures:
|
|
return ExperimentalClientFeatures(mock_session)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_poll_task_yields_until_completed(features: ExperimentalClientFeatures) -> None:
|
|
"""poll_task yields each status until terminal."""
|
|
features.get_task = make_status_sequence("working", "working", "completed") # type: ignore[method-assign]
|
|
|
|
statuses = [s.status async for s in features.poll_task("test-task")]
|
|
|
|
assert statuses == ["working", "working", "completed"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("terminal_status", ["completed", "failed", "cancelled"])
|
|
async def test_poll_task_exits_on_terminal(features: ExperimentalClientFeatures, terminal_status: TaskStatus) -> None:
|
|
"""poll_task exits immediately when task is already terminal."""
|
|
features.get_task = make_status_sequence(terminal_status) # type: ignore[method-assign]
|
|
|
|
statuses = [s.status async for s in features.poll_task("test-task")]
|
|
|
|
assert statuses == [terminal_status]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_poll_task_continues_through_input_required(features: ExperimentalClientFeatures) -> None:
|
|
"""poll_task yields input_required and continues (non-terminal)."""
|
|
features.get_task = make_status_sequence("working", "input_required", "working", "completed") # type: ignore[method-assign]
|
|
|
|
statuses = [s.status async for s in features.poll_task("test-task")]
|
|
|
|
assert statuses == ["working", "input_required", "working", "completed"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_poll_task_passes_task_id(features: ExperimentalClientFeatures) -> None:
|
|
"""poll_task passes correct task_id to get_task."""
|
|
received_ids: list[str] = []
|
|
|
|
async def mock_get_task(task_id: str) -> GetTaskResult:
|
|
received_ids.append(task_id)
|
|
return make_task_result(status="completed", task_id=task_id)
|
|
|
|
features.get_task = mock_get_task # type: ignore[method-assign]
|
|
|
|
_ = [s async for s in features.poll_task("my-task-123")]
|
|
|
|
assert received_ids == ["my-task-123"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_poll_task_yields_full_result(features: ExperimentalClientFeatures) -> None:
|
|
"""poll_task yields complete GetTaskResult objects."""
|
|
|
|
async def mock_get_task(task_id: str) -> GetTaskResult:
|
|
return make_task_result(
|
|
status="completed",
|
|
task_id=task_id,
|
|
status_message="All done!",
|
|
)
|
|
|
|
features.get_task = mock_get_task # type: ignore[method-assign]
|
|
|
|
results = [r async for r in features.poll_task("test-task")]
|
|
|
|
assert len(results) == 1
|
|
assert results[0].status == "completed"
|
|
assert results[0].status_message == "All done!"
|
|
assert results[0].task_id == "test-task"
|