forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_152_resource_mime_type.py
More file actions
146 lines (120 loc) · 5.14 KB
/
test_152_resource_mime_type.py
File metadata and controls
146 lines (120 loc) · 5.14 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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
import base64
import pytest
from pydantic import AnyUrl
from mcp import types
from mcp.server.fastmcp import FastMCP
from mcp.server.lowlevel import Server
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.shared.memory import (
create_connected_server_and_client_session as client_session,
)
pytestmark = pytest.mark.anyio
async def test_fastmcp_resource_mime_type():
"""Test that mime_type parameter is respected for resources."""
mcp = FastMCP("test")
# Create a small test image as bytes
image_bytes = b"fake_image_data"
base64_string = base64.b64encode(image_bytes).decode("utf-8")
@mcp.resource("test://image", mime_type="image/png")
def get_image_as_string() -> str:
"""Return a test image as base64 string."""
return base64_string
@mcp.resource("test://image_bytes", mime_type="image/png")
def get_image_as_bytes() -> bytes:
"""Return a test image as bytes."""
return image_bytes
# Test that resources are listed with correct mime type
async with client_session(mcp._mcp_server) as client:
# List resources and verify mime types
resources = await client.list_resources()
assert resources.resources is not None
mapping = {str(r.uri): r for r in resources.resources}
# Find our resources
string_resource = mapping["test://image"]
bytes_resource = mapping["test://image_bytes"]
# Verify mime types
assert (
string_resource.mimeType == "image/png"
), "String resource mime type not respected"
assert (
bytes_resource.mimeType == "image/png"
), "Bytes resource mime type not respected"
# Also verify the content can be read correctly
string_result = await client.read_resource(AnyUrl("test://image"))
assert len(string_result.contents) == 1
assert (
getattr(string_result.contents[0], "text") == base64_string
), "Base64 string mismatch"
assert (
string_result.contents[0].mimeType == "image/png"
), "String content mime type not preserved"
bytes_result = await client.read_resource(AnyUrl("test://image_bytes"))
assert len(bytes_result.contents) == 1
assert (
base64.b64decode(getattr(bytes_result.contents[0], "blob")) == image_bytes
), "Bytes mismatch"
assert (
bytes_result.contents[0].mimeType == "image/png"
), "Bytes content mime type not preserved"
async def test_lowlevel_resource_mime_type():
"""Test that mime_type parameter is respected for resources."""
server = Server("test")
# Create a small test image as bytes
image_bytes = b"fake_image_data"
base64_string = base64.b64encode(image_bytes).decode("utf-8")
# Create test resources with specific mime types
test_resources = [
types.Resource(
uri=AnyUrl("test://image"), name="test image", mimeType="image/png"
),
types.Resource(
uri=AnyUrl("test://image_bytes"),
name="test image bytes",
mimeType="image/png",
),
]
@server.list_resources()
async def handle_list_resources():
return test_resources
@server.read_resource()
async def handle_read_resource(uri: AnyUrl):
if str(uri) == "test://image":
return [ReadResourceContents(content=base64_string, mime_type="image/png")]
elif str(uri) == "test://image_bytes":
return [
ReadResourceContents(content=bytes(image_bytes), mime_type="image/png")
]
raise Exception(f"Resource not found: {uri}")
# Test that resources are listed with correct mime type
async with client_session(server) as client:
# List resources and verify mime types
resources = await client.list_resources()
assert resources.resources is not None
mapping = {str(r.uri): r for r in resources.resources}
# Find our resources
string_resource = mapping["test://image"]
bytes_resource = mapping["test://image_bytes"]
# Verify mime types
assert (
string_resource.mimeType == "image/png"
), "String resource mime type not respected"
assert (
bytes_resource.mimeType == "image/png"
), "Bytes resource mime type not respected"
# Also verify the content can be read correctly
string_result = await client.read_resource(AnyUrl("test://image"))
assert len(string_result.contents) == 1
assert (
getattr(string_result.contents[0], "text") == base64_string
), "Base64 string mismatch"
assert (
string_result.contents[0].mimeType == "image/png"
), "String content mime type not preserved"
bytes_result = await client.read_resource(AnyUrl("test://image_bytes"))
assert len(bytes_result.contents) == 1
assert (
base64.b64decode(getattr(bytes_result.contents[0], "blob")) == image_bytes
), "Bytes mismatch"
assert (
bytes_result.contents[0].mimeType == "image/png"
), "Bytes content mime type not preserved"