python-sdk/tests/server/test_lifespan.py at main · devilcoder/python-sdk · GitHub

History

219 lines (188 loc) · 6.81 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

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

"""Tests for lifespan functionality in both low-level and FastMCP servers."""

from collections.abc import AsyncIterator

from contextlib import asynccontextmanager

import anyio

import pytest

from pydantic import TypeAdapter

from mcp.server.fastmcp import Context, FastMCP

from mcp.server.lowlevel.server import NotificationOptions, Server

from mcp.server.models import InitializationOptions

from mcp.types import (

ClientCapabilities,

Implementation,

InitializeRequestParams,

JSONRPCMessage,

JSONRPCNotification,

JSONRPCRequest,

)

@pytest.mark.anyio

async def test_lowlevel_server_lifespan():

"""Test that lifespan works in low-level server."""

@asynccontextmanager

async def test_lifespan(server: Server) -> AsyncIterator[dict[str, bool]]:

"""Test lifespan context that tracks startup/shutdown."""

context = {"started": False, "shutdown": False}

try:

context["started"] = True

yield context

finally:

context["shutdown"] = True

server = Server("test", lifespan=test_lifespan)

# Create memory streams for testing

send_stream1, receive_stream1 = anyio.create_memory_object_stream(100)

send_stream2, receive_stream2 = anyio.create_memory_object_stream(100)

# Create a tool that accesses lifespan context

@server.call_tool()

async def check_lifespan(name: str, arguments: dict) -> list:

ctx = server.request_context

assert isinstance(ctx.lifespan_context, dict)

assert ctx.lifespan_context["started"]

assert not ctx.lifespan_context["shutdown"]

return [{"type": "text", "text": "true"}]

# Run server in background task

async with (

anyio.create_task_group() as tg,

send_stream1,

receive_stream1,

send_stream2,

receive_stream2,

async def run_server():

await server.run(

receive_stream1,

send_stream2,

InitializationOptions(

server_name="test",

server_version="0.1.0",

capabilities=server.get_capabilities(

notification_options=NotificationOptions(),

experimental_capabilities={},

raise_exceptions=True,

)

tg.start_soon(run_server)

# Initialize the server

params = InitializeRequestParams(

protocolVersion="2024-11-05",

capabilities=ClientCapabilities(),

clientInfo=Implementation(name="test-client", version="0.1.0"),

)

await send_stream1.send(

JSONRPCMessage(

root=JSONRPCRequest(

jsonrpc="2.0",

id=1,

method="initialize",

params=TypeAdapter(InitializeRequestParams).dump_python(params),

)

response = await receive_stream2.receive()

# Send initialized notification

await send_stream1.send(

JSONRPCMessage(

root=JSONRPCNotification(

jsonrpc="2.0",

method="notifications/initialized",

)

# Call the tool to verify lifespan context

await send_stream1.send(

JSONRPCMessage(

root=JSONRPCRequest(

jsonrpc="2.0",

id=2,

method="tools/call",

params={"name": "check_lifespan", "arguments": {}},

)

# Get response and verify

response = await receive_stream2.receive()

assert response.root.result["content"][0]["text"] == "true"

# Cancel server task

tg.cancel_scope.cancel()

@pytest.mark.anyio

async def test_fastmcp_server_lifespan():

"""Test that lifespan works in FastMCP server."""

@asynccontextmanager

async def test_lifespan(server: FastMCP) -> AsyncIterator[dict]:

"""Test lifespan context that tracks startup/shutdown."""

context = {"started": False, "shutdown": False}

try:

context["started"] = True

yield context

finally:

context["shutdown"] = True

server = FastMCP("test", lifespan=test_lifespan)

# Create memory streams for testing

send_stream1, receive_stream1 = anyio.create_memory_object_stream(100)

send_stream2, receive_stream2 = anyio.create_memory_object_stream(100)

# Add a tool that checks lifespan context

@server.tool()

def check_lifespan(ctx: Context) -> bool:

"""Tool that checks lifespan context."""

assert isinstance(ctx.request_context.lifespan_context, dict)

assert ctx.request_context.lifespan_context["started"]

assert not ctx.request_context.lifespan_context["shutdown"]

return True

# Run server in background task

async with (

anyio.create_task_group() as tg,

send_stream1,

receive_stream1,

send_stream2,

receive_stream2,

async def run_server():

await server._mcp_server.run(

receive_stream1,

send_stream2,

server._mcp_server.create_initialization_options(),

raise_exceptions=True,

)

tg.start_soon(run_server)

# Initialize the server

params = InitializeRequestParams(

protocolVersion="2024-11-05",

capabilities=ClientCapabilities(),

clientInfo=Implementation(name="test-client", version="0.1.0"),

)

await send_stream1.send(

JSONRPCMessage(

root=JSONRPCRequest(

jsonrpc="2.0",

id=1,

method="initialize",

params=TypeAdapter(InitializeRequestParams).dump_python(params),

)

response = await receive_stream2.receive()

# Send initialized notification

await send_stream1.send(

JSONRPCMessage(

root=JSONRPCNotification(

jsonrpc="2.0",

method="notifications/initialized",

)

# Call the tool to verify lifespan context

await send_stream1.send(

JSONRPCMessage(

root=JSONRPCRequest(

jsonrpc="2.0",

id=2,

method="tools/call",

params={"name": "check_lifespan", "arguments": {}},

)

# Get response and verify

response = await receive_stream2.receive()

assert response.root.result["content"][0]["text"] == "true"

# Cancel server task

tg.cancel_scope.cancel()

python-sdk/tests/server/test_lifespan.py at main · devilcoder/python-sdk · GitHub

Expand file tree

Breadcrumbs

test_lifespan.py

Latest commit

History

Breadcrumbs

test_lifespan.py

File metadata and controls