Skip to content

Commit b529e26

Browse files
maxisbeymaxisbey
committed
Achieve 100% coverage for experimental tasks
Test improvements: - Replace unreachable handler bodies with raise NotImplementedError - Use pragma: no branch for no-op message handlers (ellipsis bodies) - Remove dead code (unused helper functions) - Add try/finally blocks for stream cleanup with pragma: no cover - Simplify handler error paths to use assert instead of if/return - Remove unnecessary store.cleanup() calls after cancelled task groups Source improvements: - Add pragma: no cover for defensive _meta checks (unreachable code paths)
1 parent 49543bb commit b529e26

File tree

8 files changed

+78
-189
lines changed

8 files changed

+78
-189
lines changed

‎tests/experimental/tasks/client/test_capabilities.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -92,18 +92,18 @@ async def test_client_capabilities_with_tasks():
9292

9393
received_capabilities: ClientCapabilities | None = None
9494

95-
# Define custom handlers to trigger capability building
95+
# Define custom handlers to trigger capability building (never actually called)
9696
async def my_list_tasks_handler(
9797
context: RequestContext[ClientSession, None],
9898
params: types.PaginatedRequestParams | None,
9999
) -> types.ListTasksResult | types.ErrorData:
100-
return types.ListTasksResult(tasks=[])
100+
raise NotImplementedError
101101

102102
async def my_cancel_task_handler(
103103
context: RequestContext[ClientSession, None],
104104
params: types.CancelTaskRequestParams,
105105
) -> types.CancelTaskResult | types.ErrorData:
106-
return types.ErrorData(code=types.INVALID_REQUEST, message="Not found")
106+
raise NotImplementedError
107107

108108
async def mock_server():
109109
nonlocal received_capabilities
@@ -181,13 +181,13 @@ async def my_list_tasks_handler(
181181
context: RequestContext[ClientSession, None],
182182
params: types.PaginatedRequestParams | None,
183183
) -> types.ListTasksResult | types.ErrorData:
184-
return types.ListTasksResult(tasks=[])
184+
raise NotImplementedError
185185

186186
async def my_cancel_task_handler(
187187
context: RequestContext[ClientSession, None],
188188
params: types.CancelTaskRequestParams,
189189
) -> types.CancelTaskResult | types.ErrorData:
190-
return types.ErrorData(code=types.INVALID_REQUEST, message="Not found")
190+
raise NotImplementedError
191191

192192
async def mock_server():
193193
nonlocal received_capabilities
@@ -267,7 +267,7 @@ async def my_augmented_sampling_handler(
267267
params: types.CreateMessageRequestParams,
268268
task_metadata: types.TaskMetadata,
269269
) -> types.CreateTaskResult | types.ErrorData:
270-
return types.ErrorData(code=types.INVALID_REQUEST, message="Not implemented")
270+
raise NotImplementedError
271271

272272
async def mock_server():
273273
nonlocal received_capabilities

‎tests/experimental/tasks/client/test_handlers.py‎

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -101,9 +101,8 @@ async def client_streams() -> AsyncIterator[ClientTestStreams]:
101101
async def _default_message_handler(
102102
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
103103
) -> None:
104-
"""Default message handler that re-raises exceptions."""
105-
if isinstance(message, Exception):
106-
raise message
104+
"""Default message handler that ignores messages (tests handle them explicitly)."""
105+
...
107106

108107

109108
@pytest.mark.anyio
@@ -120,8 +119,7 @@ async def get_task_handler(
120119
nonlocal received_task_id
121120
received_task_id = params.taskId
122121
task = await store.get_task(params.taskId)
123-
if task is None:
124-
return ErrorData(code=types.INVALID_REQUEST, message=f"Task {params.taskId} not found")
122+
assert task is not None, f"Test setup error: task {params.taskId} should exist"
125123
return GetTaskResult(
126124
taskId=task.taskId,
127125
status=task.status,
@@ -186,8 +184,7 @@ async def get_task_result_handler(
186184
params: GetTaskPayloadRequestParams,
187185
) -> GetTaskPayloadResult | ErrorData:
188186
result = await store.get_result(params.taskId)
189-
if result is None:
190-
return ErrorData(code=types.INVALID_REQUEST, message=f"Result for {params.taskId} not found")
187+
assert result is not None, f"Test setup error: result for {params.taskId} should exist"
191188
assert isinstance(result, types.CallToolResult)
192189
return GetTaskPayloadResult(**result.model_dump())
193190

@@ -305,8 +302,7 @@ async def cancel_task_handler(
305302
params: CancelTaskRequestParams,
306303
) -> CancelTaskResult | ErrorData:
307304
task = await store.get_task(params.taskId)
308-
if task is None:
309-
return ErrorData(code=types.INVALID_REQUEST, message=f"Task {params.taskId} not found")
305+
assert task is not None, f"Test setup error: task {params.taskId} should exist"
310306
await store.update_task(params.taskId, status="cancelled")
311307
updated = await store.get_task(params.taskId)
312308
assert updated is not None
@@ -396,8 +392,7 @@ async def get_task_handler(
396392
params: GetTaskRequestParams,
397393
) -> GetTaskResult | ErrorData:
398394
task = await store.get_task(params.taskId)
399-
if task is None:
400-
return ErrorData(code=types.INVALID_REQUEST, message="Task not found")
395+
assert task is not None, f"Test setup error: task {params.taskId} should exist"
401396
return GetTaskResult(
402397
taskId=task.taskId,
403398
status=task.status,
@@ -413,8 +408,7 @@ async def get_task_result_handler(
413408
params: GetTaskPayloadRequestParams,
414409
) -> GetTaskPayloadResult | ErrorData:
415410
result = await store.get_result(params.taskId)
416-
if result is None:
417-
return ErrorData(code=types.INVALID_REQUEST, message="Result not found")
411+
assert result is not None, f"Test setup error: result for {params.taskId} should exist"
418412
assert isinstance(result, CreateMessageResult)
419413
return GetTaskPayloadResult(**result.model_dump())
420414

@@ -538,8 +532,7 @@ async def get_task_handler(
538532
params: GetTaskRequestParams,
539533
) -> GetTaskResult | ErrorData:
540534
task = await store.get_task(params.taskId)
541-
if task is None:
542-
return ErrorData(code=types.INVALID_REQUEST, message="Task not found")
535+
assert task is not None, f"Test setup error: task {params.taskId} should exist"
543536
return GetTaskResult(
544537
taskId=task.taskId,
545538
status=task.status,
@@ -555,8 +548,7 @@ async def get_task_result_handler(
555548
params: GetTaskPayloadRequestParams,
556549
) -> GetTaskPayloadResult | ErrorData:
557550
result = await store.get_result(params.taskId)
558-
if result is None:
559-
return ErrorData(code=types.INVALID_REQUEST, message="Result not found")
551+
assert result is not None, f"Test setup error: result for {params.taskId} should exist"
560552
assert isinstance(result, ElicitResult)
561553
return GetTaskPayloadResult(**result.model_dump())
562554

‎tests/experimental/tasks/client/test_tasks.py‎

Lines changed: 15 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -81,14 +81,13 @@ async def do_work():
8181
app.task_group.start_soon(do_work)
8282
return CreateTaskResult(task=task)
8383

84-
return [TextContent(type="text", text="Sync")]
84+
raise NotImplementedError
8585

8686
@server.experimental.get_task()
8787
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
8888
app = server.request_context.lifespan_context
8989
task = await app.store.get_task(request.params.taskId)
90-
if task is None:
91-
raise ValueError(f"Task {request.params.taskId} not found")
90+
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
9291
return GetTaskResult(
9392
taskId=task.taskId,
9493
status=task.status,
@@ -105,9 +104,7 @@ async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
105104

106105
async def message_handler(
107106
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
108-
) -> None:
109-
if isinstance(message, Exception):
110-
raise message
107+
) -> None: ... # pragma: no branch
111108

112109
async def run_server(app_context: AppContext):
113110
async with ServerSession(
@@ -162,8 +159,6 @@ async def run_server(app_context: AppContext):
162159

163160
tg.cancel_scope.cancel()
164161

165-
store.cleanup()
166-
167162

168163
@pytest.mark.anyio
169164
async def test_session_experimental_get_task_result() -> None:
@@ -198,14 +193,15 @@ async def do_work():
198193
app.task_group.start_soon(do_work)
199194
return CreateTaskResult(task=task)
200195

201-
return [TextContent(type="text", text="Sync")]
196+
raise NotImplementedError
202197

203198
@server.experimental.get_task_result()
204-
async def handle_get_task_result(request: GetTaskPayloadRequest) -> GetTaskPayloadResult:
199+
async def handle_get_task_result(
200+
request: GetTaskPayloadRequest,
201+
) -> GetTaskPayloadResult:
205202
app = server.request_context.lifespan_context
206203
result = await app.store.get_result(request.params.taskId)
207-
if result is None:
208-
raise ValueError(f"Result for task {request.params.taskId} not found")
204+
assert result is not None, f"Test setup error: result for {request.params.taskId} should exist"
209205
assert isinstance(result, CallToolResult)
210206
return GetTaskPayloadResult(**result.model_dump())
211207

@@ -215,9 +211,7 @@ async def handle_get_task_result(request: GetTaskPayloadRequest) -> GetTaskPaylo
215211

216212
async def message_handler(
217213
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
218-
) -> None:
219-
if isinstance(message, Exception):
220-
raise message
214+
) -> None: ... # pragma: no branch
221215

222216
async def run_server(app_context: AppContext):
223217
async with ServerSession(
@@ -274,8 +268,6 @@ async def run_server(app_context: AppContext):
274268

275269
tg.cancel_scope.cancel()
276270

277-
store.cleanup()
278-
279271

280272
@pytest.mark.anyio
281273
async def test_session_experimental_list_tasks() -> None:
@@ -308,7 +300,7 @@ async def do_work():
308300
app.task_group.start_soon(do_work)
309301
return CreateTaskResult(task=task)
310302

311-
return [TextContent(type="text", text="Sync")]
303+
raise NotImplementedError
312304

313305
@server.experimental.list_tasks()
314306
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
@@ -322,9 +314,7 @@ async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
322314

323315
async def message_handler(
324316
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
325-
) -> None:
326-
if isinstance(message, Exception):
327-
raise message
317+
) -> None: ... # pragma: no branch
328318

329319
async def run_server(app_context: AppContext):
330320
async with ServerSession(
@@ -376,8 +366,6 @@ async def run_server(app_context: AppContext):
376366

377367
tg.cancel_scope.cancel()
378368

379-
store.cleanup()
380-
381369

382370
@pytest.mark.anyio
383371
async def test_session_experimental_cancel_task() -> None:
@@ -401,14 +389,13 @@ async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextCon
401389
# Don't start any work - task stays in "working" status
402390
return CreateTaskResult(task=task)
403391

404-
return [TextContent(type="text", text="Sync")]
392+
raise NotImplementedError
405393

406394
@server.experimental.get_task()
407395
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
408396
app = server.request_context.lifespan_context
409397
task = await app.store.get_task(request.params.taskId)
410-
if task is None:
411-
raise ValueError(f"Task {request.params.taskId} not found")
398+
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
412399
return GetTaskResult(
413400
taskId=task.taskId,
414401
status=task.status,
@@ -423,8 +410,7 @@ async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
423410
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
424411
app = server.request_context.lifespan_context
425412
task = await app.store.get_task(request.params.taskId)
426-
if task is None:
427-
raise ValueError(f"Task {request.params.taskId} not found")
413+
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
428414
await app.store.update_task(request.params.taskId, status="cancelled")
429415
# CancelTaskResult extends Task, so we need to return the updated task info
430416
updated_task = await app.store.get_task(request.params.taskId)
@@ -443,9 +429,7 @@ async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
443429

444430
async def message_handler(
445431
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
446-
) -> None:
447-
if isinstance(message, Exception):
448-
raise message
432+
) -> None: ... # pragma: no branch
449433

450434
async def run_server(app_context: AppContext):
451435
async with ServerSession(
@@ -501,5 +485,3 @@ async def run_server(app_context: AppContext):
501485
assert status_after.status == "cancelled"
502486

503487
tg.cancel_scope.cancel()
504-
505-
store.cleanup()

‎tests/experimental/tasks/server/test_context.py‎

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,12 @@
11
"""Tests for TaskContext and helper functions."""
22

3-
import anyio
43
import pytest
54

65
from mcp.shared.experimental.tasks.context import TaskContext
76
from mcp.shared.experimental.tasks.helpers import create_task_state, task_execution
87
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
98
from mcp.types import CallToolResult, TaskMetadata, TextContent
109

11-
12-
async def wait_for_terminal_status(store: InMemoryTaskStore, task_id: str, timeout: float = 5.0) -> None:
13-
"""Wait for a task to reach terminal status (completed, failed, cancelled)."""
14-
terminal_statuses = {"completed", "failed", "cancelled"}
15-
with anyio.fail_after(timeout):
16-
while True:
17-
task = await store.get_task(task_id)
18-
if task and task.status in terminal_statuses:
19-
return
20-
await anyio.sleep(0) # Yield to allow other tasks to run
21-
22-
2310
# --- TaskContext tests ---
2411

2512

@@ -201,6 +188,4 @@ async def test_task_execution_not_found() -> None:
201188

202189
with pytest.raises(ValueError, match="not found"):
203190
async with task_execution("nonexistent", store):
204-
pass
205-
206-
store.cleanup()
191+
...

0 commit comments

Comments
 (0)