python-sdk/tests/shared/test_context.py at main · modelcontextprotocol/python-sdk · GitHub

History

133 lines (103 loc) · 5.24 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

"""Tests for `BaseContext`.

`BaseContext` is composition over a `DispatchContext` - it forwards

`transport`/`cancel_requested`/`send_raw_request`/`notify`/`progress`

and adds `meta`. It must satisfy `Outbound` so `ClientPeer` can wrap it.

"""

from collections.abc import Mapping

from typing import Any

import anyio

import pytest

from mcp.shared.context import BaseContext

from mcp.shared.dispatcher import DispatchContext

from mcp.shared.peer import ClientPeer

from mcp.shared.transport_context import TransportContext

from .conftest import direct_pair, jsonrpc_pair

from .test_dispatcher import Recorder, echo_handlers, running_pair

DCtx = DispatchContext[TransportContext]

@pytest.mark.anyio

async def test_base_context_forwards_transport_and_cancel_requested():

captured: list[BaseContext[TransportContext]] = []

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

captured.append(bctx)

return {}

async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)

bctx = captured[0]

assert bctx.transport.kind == "direct"

assert isinstance(bctx.cancel_requested, anyio.Event)

assert bctx.can_send_request is True

assert bctx.meta is None

@pytest.mark.anyio

async def test_base_context_can_send_request_reflects_dispatch_context_closed_state():

"""`can_send_request` must track the dctx, not the static transport flag,

so it agrees with whether `send_raw_request` would raise."""

captured: list[BaseContext[TransportContext]] = []

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

captured.append(BaseContext(ctx))

return {}

async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)

bctx = captured[0]

assert bctx.transport.can_send_request is True

assert bctx.can_send_request is False

@pytest.mark.anyio

async def test_base_context_send_raw_request_and_notify_forward_to_dispatch_context():

crec = Recorder()

c_req, c_notify = echo_handlers(crec)

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

sample = await bctx.send_raw_request("sampling/createMessage", {"x": 1})

await bctx.notify("notifications/message", {"level": "info"})

return {"sample": sample}

async with running_pair(

direct_pair,

server_on_request=server_on_request,

client_on_request=c_req,

client_on_notify=c_notify,

) as (client, *_):

with anyio.fail_after(5):

result = await client.send_raw_request("tools/call", None)

await crec.notified.wait()

assert crec.requests == [("sampling/createMessage", {"x": 1})]

assert crec.notifications == [("notifications/message", {"level": "info"})]

assert result["sample"] == {"echoed": "sampling/createMessage", "params": {"x": 1}}

@pytest.mark.anyio

async def test_base_context_report_progress_invokes_caller_on_progress():

received: list[tuple[float, float | None, str | None]] = []

async def on_progress(progress: float, total: float | None, message: str | None) -> None:

received.append((progress, total, message))

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

await bctx.report_progress(0.5, total=1.0, message="halfway")

return {}

async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None, {"on_progress": on_progress})

assert received == [(0.5, 1.0, "halfway")]

@pytest.mark.anyio

async def test_base_context_satisfies_outbound_so_peer_mixin_works():

"""Wrapping a BaseContext in ClientPeer proves it satisfies Outbound structurally."""

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx)

await ClientPeer(bctx).ping()

return {}

crec = Recorder()

c_req, c_notify = echo_handlers(crec)

async with running_pair(

direct_pair, server_on_request=server_on_request, client_on_request=c_req, client_on_notify=c_notify

) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)

assert crec.requests == [("ping", None)]

@pytest.mark.anyio

async def test_base_context_meta_holds_supplied_request_params_meta():

async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:

bctx = BaseContext(ctx, meta={"progressToken": "abc"})

assert bctx.meta is not None and bctx.meta.get("progressToken") == "abc"

return {}

async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):

with anyio.fail_after(5):

await client.send_raw_request("t", None)

python-sdk/tests/shared/test_context.py at main · modelcontextprotocol/python-sdk · GitHub

Uh oh!

Expand file tree

Breadcrumbs

test_context.py

Latest commit

History

Breadcrumbs

test_context.py

File metadata and controls