mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-13 00:59:10 +02:00
654 lines
22 KiB
Python
654 lines
22 KiB
Python
from pathlib import Path
|
|
from tempfile import NamedTemporaryFile
|
|
|
|
import pytest
|
|
from pydantic import AnyUrl, FileUrl
|
|
|
|
from fastmcp.exceptions import NotFoundError, ResourceError
|
|
from fastmcp.resources import (
|
|
FileResource,
|
|
ResourceManager,
|
|
ResourceTemplate,
|
|
)
|
|
from fastmcp.resources.resource import FunctionResource
|
|
from fastmcp.utilities.tests import caplog_for_fastmcp
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_file():
|
|
"""Create a temporary file for testing.
|
|
|
|
File is automatically cleaned up after the test if it still exists.
|
|
"""
|
|
content = "test content"
|
|
with NamedTemporaryFile(mode="w", delete=False) as f:
|
|
f.write(content)
|
|
path = Path(f.name).resolve()
|
|
yield path
|
|
try:
|
|
path.unlink()
|
|
except FileNotFoundError:
|
|
pass # File was already deleted by the test
|
|
|
|
|
|
class TestResourceManager:
|
|
"""Test ResourceManager functionality."""
|
|
|
|
async def test_add_resource(self, temp_file: Path):
|
|
"""Test adding a resource."""
|
|
manager = ResourceManager()
|
|
file_url = "file://test-resource"
|
|
resource = FileResource(
|
|
uri=FileUrl(file_url),
|
|
name="test",
|
|
path=temp_file,
|
|
)
|
|
added = manager.add_resource(resource)
|
|
assert added == resource
|
|
# Get the actual key from the resource manager
|
|
resources = await manager.get_resources()
|
|
assert len(resources) == 1
|
|
assert resource in resources.values()
|
|
|
|
async def test_add_duplicate_resource(self, temp_file: Path):
|
|
"""Test adding the same resource twice."""
|
|
manager = ResourceManager()
|
|
file_url = "file://test-resource"
|
|
resource = FileResource(
|
|
uri=FileUrl(file_url),
|
|
name="test",
|
|
path=temp_file,
|
|
)
|
|
first = manager.add_resource(resource)
|
|
second = manager.add_resource(resource)
|
|
assert first == second
|
|
# Check the resource is there
|
|
resources = await manager.get_resources()
|
|
assert len(resources) == 1
|
|
assert resource in resources.values()
|
|
|
|
async def test_warn_on_duplicate_resources(self, temp_file: Path, caplog):
|
|
"""Test warning on duplicate resources."""
|
|
manager = ResourceManager(duplicate_behavior="warn")
|
|
|
|
file_url = "file://test-resource"
|
|
resource = FileResource(
|
|
uri=FileUrl(file_url),
|
|
name="test_resource",
|
|
path=temp_file,
|
|
)
|
|
|
|
manager.add_resource(resource)
|
|
|
|
with caplog_for_fastmcp(caplog):
|
|
manager.add_resource(resource)
|
|
|
|
assert "Resource already exists" in caplog.text
|
|
# Should have the resource
|
|
resources = await manager.get_resources()
|
|
assert len(resources) == 1
|
|
assert resource in resources.values()
|
|
|
|
async def test_disable_warn_on_duplicate_resources(self, temp_file: Path, caplog):
|
|
"""Test disabling warning on duplicate resources."""
|
|
manager = ResourceManager(duplicate_behavior="ignore")
|
|
resource = FileResource(
|
|
uri=FileUrl(f"file://{temp_file.name}"),
|
|
name="test",
|
|
path=temp_file,
|
|
)
|
|
manager.add_resource(resource)
|
|
manager.add_resource(resource)
|
|
assert "Resource already exists" not in caplog.text
|
|
|
|
async def test_error_on_duplicate_resources(self, temp_file: Path):
|
|
"""Test error on duplicate resources."""
|
|
manager = ResourceManager(duplicate_behavior="error")
|
|
|
|
resource = FileResource(
|
|
uri=FileUrl(f"file://{temp_file.name}"),
|
|
name="test_resource",
|
|
path=temp_file,
|
|
)
|
|
|
|
manager.add_resource(resource)
|
|
|
|
with pytest.raises(ValueError, match="Resource already exists"):
|
|
manager.add_resource(resource)
|
|
|
|
async def test_replace_duplicate_resources(self, temp_file: Path):
|
|
"""Test replacing duplicate resources."""
|
|
manager = ResourceManager(duplicate_behavior="replace")
|
|
|
|
file_url = "file://test-resource"
|
|
resource1 = FileResource(
|
|
uri=FileUrl(file_url),
|
|
name="original",
|
|
path=temp_file,
|
|
)
|
|
|
|
resource2 = FileResource(
|
|
uri=FileUrl(file_url),
|
|
name="replacement",
|
|
path=temp_file,
|
|
)
|
|
|
|
manager.add_resource(resource1)
|
|
manager.add_resource(resource2)
|
|
|
|
# Should have replaced with the new resource
|
|
resources = await manager.get_resources()
|
|
resource_list = list(resources.values())
|
|
assert len(resource_list) == 1
|
|
assert resource_list[0].name == "replacement"
|
|
|
|
async def test_ignore_duplicate_resources(self, temp_file: Path):
|
|
"""Test ignoring duplicate resources."""
|
|
manager = ResourceManager(duplicate_behavior="ignore")
|
|
|
|
file_url = "file://test-resource"
|
|
resource1 = FileResource(
|
|
uri=FileUrl(file_url),
|
|
name="original",
|
|
path=temp_file,
|
|
)
|
|
|
|
resource2 = FileResource(
|
|
uri=FileUrl(file_url),
|
|
name="replacement",
|
|
path=temp_file,
|
|
)
|
|
|
|
manager.add_resource(resource1)
|
|
result = manager.add_resource(resource2)
|
|
|
|
# Should keep the original
|
|
resources = await manager.get_resources()
|
|
resource_list = list(resources.values())
|
|
assert len(resource_list) == 1
|
|
assert resource_list[0].name == "original"
|
|
# Result should be the original resource
|
|
assert result.name == "original"
|
|
|
|
async def test_warn_on_duplicate_templates(self, caplog):
|
|
"""Test warning on duplicate templates."""
|
|
manager = ResourceManager(duplicate_behavior="warn")
|
|
|
|
def template_fn(id: str) -> str:
|
|
return f"Template {id}"
|
|
|
|
template = ResourceTemplate.from_function(
|
|
fn=template_fn,
|
|
uri_template="test://{id}",
|
|
name="test_template",
|
|
)
|
|
|
|
manager.add_template(template)
|
|
|
|
with caplog_for_fastmcp(caplog):
|
|
manager.add_template(template)
|
|
|
|
assert "Template already exists" in caplog.text
|
|
# Should have the template
|
|
templates = await manager.get_resource_templates()
|
|
assert templates == {"test://{id}": template}
|
|
|
|
async def test_error_on_duplicate_templates(self):
|
|
"""Test error on duplicate templates."""
|
|
manager = ResourceManager(duplicate_behavior="error")
|
|
|
|
def template_fn(id: str) -> str:
|
|
return f"Template {id}"
|
|
|
|
template = ResourceTemplate.from_function(
|
|
fn=template_fn,
|
|
uri_template="test://{id}",
|
|
name="test_template",
|
|
)
|
|
|
|
manager.add_template(template)
|
|
|
|
with pytest.raises(ValueError, match="Template already exists"):
|
|
manager.add_template(template)
|
|
|
|
async def test_replace_duplicate_templates(self):
|
|
"""Test replacing duplicate templates."""
|
|
manager = ResourceManager(duplicate_behavior="replace")
|
|
|
|
def original_fn(id: str) -> str:
|
|
return f"Original {id}"
|
|
|
|
def replacement_fn(id: str) -> str:
|
|
return f"Replacement {id}"
|
|
|
|
template1 = ResourceTemplate.from_function(
|
|
fn=original_fn,
|
|
uri_template="test://{id}",
|
|
name="original",
|
|
)
|
|
|
|
template2 = ResourceTemplate.from_function(
|
|
fn=replacement_fn,
|
|
uri_template="test://{id}",
|
|
name="replacement",
|
|
)
|
|
|
|
manager.add_template(template1)
|
|
manager.add_template(template2)
|
|
|
|
# Should have replaced with the new template
|
|
templates_dict = await manager.get_resource_templates()
|
|
templates = list(templates_dict.values())
|
|
assert len(templates) == 1
|
|
assert templates[0].name == "replacement"
|
|
|
|
async def test_ignore_duplicate_templates(self):
|
|
"""Test ignoring duplicate templates."""
|
|
manager = ResourceManager(duplicate_behavior="ignore")
|
|
|
|
def original_fn(id: str) -> str:
|
|
return f"Original {id}"
|
|
|
|
def replacement_fn(id: str) -> str:
|
|
return f"Replacement {id}"
|
|
|
|
template1 = ResourceTemplate.from_function(
|
|
fn=original_fn,
|
|
uri_template="test://{id}",
|
|
name="original",
|
|
)
|
|
|
|
template2 = ResourceTemplate.from_function(
|
|
fn=replacement_fn,
|
|
uri_template="test://{id}",
|
|
name="replacement",
|
|
)
|
|
|
|
manager.add_template(template1)
|
|
result = manager.add_template(template2)
|
|
|
|
# Should keep the original
|
|
templates_dict = await manager.get_resource_templates()
|
|
templates = list(templates_dict.values())
|
|
assert len(templates) == 1
|
|
assert templates[0].name == "original"
|
|
# Result should be the original template
|
|
assert result.name == "original"
|
|
|
|
async def test_get_resource(self, temp_file: Path):
|
|
"""Test getting a resource by URI."""
|
|
manager = ResourceManager()
|
|
resource = FileResource(
|
|
uri=FileUrl(f"file://{temp_file.name}"),
|
|
name="test",
|
|
path=temp_file,
|
|
)
|
|
manager.add_resource(resource)
|
|
retrieved = await manager.get_resource(resource.uri)
|
|
assert retrieved == resource
|
|
|
|
async def test_get_resource_from_template(self):
|
|
"""Test getting a resource through a template."""
|
|
manager = ResourceManager()
|
|
|
|
def greet(name: str) -> str:
|
|
return f"Hello, {name}!"
|
|
|
|
template = ResourceTemplate.from_function(
|
|
fn=greet,
|
|
uri_template="greet://{name}",
|
|
name="greeter",
|
|
)
|
|
manager._templates[template.uri_template] = template
|
|
|
|
resource = await manager.get_resource(AnyUrl("greet://world"))
|
|
assert isinstance(resource, FunctionResource)
|
|
content = await resource.read()
|
|
assert content == "Hello, world!"
|
|
|
|
async def test_get_unknown_resource(self):
|
|
"""Test getting a non-existent resource."""
|
|
manager = ResourceManager()
|
|
with pytest.raises(NotFoundError, match="Unknown resource"):
|
|
await manager.get_resource(AnyUrl("unknown://test"))
|
|
|
|
async def test_get_resources(self, temp_file: Path):
|
|
"""Test retrieving all resources."""
|
|
manager = ResourceManager()
|
|
file_url1 = "file://test-resource1"
|
|
resource1 = FileResource(
|
|
uri=FileUrl(file_url1),
|
|
name="test1",
|
|
path=temp_file,
|
|
)
|
|
file_url2 = "file://test-resource2"
|
|
resource2 = FileResource(
|
|
uri=FileUrl(file_url2),
|
|
name="test2",
|
|
path=temp_file,
|
|
)
|
|
manager.add_resource(resource1)
|
|
manager.add_resource(resource2)
|
|
resources = await manager.get_resources()
|
|
assert len(resources) == 2
|
|
values = list(resources.values())
|
|
assert resource1 in values
|
|
assert resource2 in values
|
|
|
|
|
|
class TestResourceTags:
|
|
"""Test functionality related to resource tags."""
|
|
|
|
async def test_add_resource_with_tags(self, temp_file: Path):
|
|
"""Test adding a resource with tags."""
|
|
manager = ResourceManager()
|
|
resource = FileResource(
|
|
uri=FileUrl("file://weather-data"),
|
|
name="weather_data",
|
|
path=temp_file,
|
|
tags={"weather", "data"},
|
|
)
|
|
manager.add_resource(resource)
|
|
|
|
# Check that tags are preserved
|
|
resources_dict = await manager.get_resources()
|
|
resources = list(resources_dict.values())
|
|
assert len(resources) == 1
|
|
assert resources[0].tags == {"weather", "data"}
|
|
|
|
async def test_add_function_resource_with_tags(self):
|
|
"""Test adding a function resource with tags."""
|
|
manager = ResourceManager()
|
|
|
|
async def get_data():
|
|
return "Sample data"
|
|
|
|
resource = FunctionResource(
|
|
uri=AnyUrl("data://sample"),
|
|
name="sample_data",
|
|
description="Sample data resource",
|
|
mime_type="text/plain",
|
|
fn=get_data,
|
|
tags={"sample", "test", "data"},
|
|
)
|
|
|
|
manager.add_resource(resource)
|
|
resources_dict = await manager.get_resources()
|
|
resources = list(resources_dict.values())
|
|
assert len(resources) == 1
|
|
assert resources[0].tags == {"sample", "test", "data"}
|
|
|
|
async def test_add_template_with_tags(self):
|
|
"""Test adding a resource template with tags."""
|
|
manager = ResourceManager()
|
|
|
|
def user_data(user_id: str) -> str:
|
|
return f"Data for user {user_id}"
|
|
|
|
template = ResourceTemplate.from_function(
|
|
fn=user_data,
|
|
uri_template="users://{user_id}",
|
|
name="user_template",
|
|
description="Get user data by ID",
|
|
tags={"users", "template", "data"},
|
|
)
|
|
|
|
manager.add_template(template)
|
|
templates_dict = await manager.get_resource_templates()
|
|
templates = list(templates_dict.values())
|
|
assert len(templates) == 1
|
|
assert templates[0].tags == {"users", "template", "data"}
|
|
|
|
async def test_filter_resources_by_tags(self, temp_file: Path):
|
|
"""Test filtering resources by tags."""
|
|
manager = ResourceManager()
|
|
|
|
# Create multiple resources with different tags
|
|
resource1 = FileResource(
|
|
uri=FileUrl("file://weather-data"),
|
|
name="weather_data",
|
|
path=temp_file,
|
|
tags={"weather", "data"},
|
|
)
|
|
|
|
async def get_user_data():
|
|
return "User data"
|
|
|
|
resource2 = FunctionResource(
|
|
uri=AnyUrl("data://users"),
|
|
name="user_data",
|
|
description="User data resource",
|
|
mime_type="text/plain",
|
|
fn=get_user_data,
|
|
tags={"users", "data"},
|
|
)
|
|
|
|
async def get_system_data():
|
|
return "System data"
|
|
|
|
resource3 = FunctionResource(
|
|
uri=AnyUrl("data://system"),
|
|
name="system_data",
|
|
description="System data resource",
|
|
mime_type="text/plain",
|
|
fn=get_system_data,
|
|
tags={"system", "admin"},
|
|
)
|
|
|
|
manager.add_resource(resource1)
|
|
manager.add_resource(resource2)
|
|
manager.add_resource(resource3)
|
|
|
|
# Filter by tags
|
|
resources_dict = await manager.get_resources()
|
|
data_resources = [r for r in resources_dict.values() if "data" in r.tags]
|
|
assert len(data_resources) == 2
|
|
assert {r.name for r in data_resources} == {"weather_data", "user_data"}
|
|
|
|
admin_resources = [r for r in resources_dict.values() if "admin" in r.tags]
|
|
assert len(admin_resources) == 1
|
|
assert admin_resources[0].name == "system_data"
|
|
|
|
|
|
class TestCustomResourceKeys:
|
|
"""Test adding resources and templates with custom keys."""
|
|
|
|
def test_add_resource_with_custom_key(self, temp_file: Path):
|
|
"""Test adding a resource with a custom key different from its URI."""
|
|
manager = ResourceManager()
|
|
original_uri = "data://test/resource"
|
|
custom_key = "custom://resource/key"
|
|
|
|
# Create a function resource instead of file resource to avoid path issues
|
|
async def get_data():
|
|
return "Test data"
|
|
|
|
resource = FunctionResource(
|
|
uri=AnyUrl(original_uri),
|
|
name="test_resource",
|
|
fn=get_data,
|
|
)
|
|
|
|
# Use model_copy to create a new resource with the custom key
|
|
resource_with_custom_key = resource.model_copy(key=custom_key)
|
|
manager.add_resource(resource_with_custom_key)
|
|
|
|
# Resource should be accessible via custom key
|
|
assert custom_key in manager._resources
|
|
# But not via its original URI
|
|
assert original_uri not in manager._resources
|
|
# The resource's internal URI remains unchanged
|
|
assert str(manager._resources[custom_key].uri) == original_uri
|
|
|
|
def test_add_template_with_custom_key(self):
|
|
"""Test adding a template with a custom key different from its URI template."""
|
|
manager = ResourceManager()
|
|
|
|
def template_fn(id: str) -> str:
|
|
return f"Template {id}"
|
|
|
|
original_uri_template = "test://{id}"
|
|
custom_key = "custom://{id}/template"
|
|
|
|
template = ResourceTemplate.from_function(
|
|
fn=template_fn,
|
|
uri_template=original_uri_template,
|
|
name="test_template",
|
|
)
|
|
|
|
# Use model_copy to create a new template with the custom key
|
|
template_with_custom_key = template.model_copy(key=custom_key)
|
|
manager.add_template(template_with_custom_key)
|
|
|
|
# Template should be accessible via custom key
|
|
assert custom_key in manager._templates
|
|
# But not via its original URI template
|
|
assert original_uri_template not in manager._templates
|
|
# The template's internal URI template remains unchanged
|
|
assert str(manager._templates[custom_key].uri_template) == original_uri_template
|
|
|
|
async def test_get_resource_with_custom_key(self, temp_file: Path):
|
|
"""Test that get_resource works with resources added with custom keys."""
|
|
manager = ResourceManager()
|
|
original_uri = "data://test/resource"
|
|
custom_key = "custom://resource/path"
|
|
|
|
# Create a function resource instead of file resource to avoid path issues
|
|
async def get_data():
|
|
return "Test data"
|
|
|
|
resource = FunctionResource(
|
|
uri=AnyUrl(original_uri),
|
|
name="test_resource",
|
|
fn=get_data,
|
|
)
|
|
|
|
# Use model_copy to create a new resource with the custom key
|
|
resource_with_custom_key = resource.model_copy(key=custom_key)
|
|
manager.add_resource(resource_with_custom_key)
|
|
|
|
# Should be retrievable by the custom key
|
|
retrieved = await manager.get_resource(custom_key)
|
|
assert retrieved is not None
|
|
assert str(retrieved.uri) == original_uri
|
|
|
|
# Should NOT be retrievable by the original URI
|
|
with pytest.raises(NotFoundError, match="Unknown resource"):
|
|
await manager.get_resource(original_uri)
|
|
|
|
async def test_get_resource_from_template_with_custom_key(self):
|
|
"""Test that templates with custom keys can create resources."""
|
|
manager = ResourceManager()
|
|
|
|
def greet(name: str) -> str:
|
|
return f"Hello, {name}!"
|
|
|
|
original_template = "greet://{name}"
|
|
custom_key = "custom://greet/{name}"
|
|
|
|
template = ResourceTemplate.from_function(
|
|
fn=greet,
|
|
uri_template=original_template,
|
|
name="custom_greeter",
|
|
)
|
|
|
|
# Use model_copy to create a new template with the custom key
|
|
template_with_custom_key = template.model_copy(key=custom_key)
|
|
manager.add_template(template_with_custom_key)
|
|
|
|
# Using a URI that matches the custom key pattern
|
|
resource = await manager.get_resource("custom://greet/world")
|
|
assert isinstance(resource, FunctionResource)
|
|
content = await resource.read()
|
|
assert content == "Hello, world!"
|
|
|
|
# Shouldn't work with the original template pattern
|
|
with pytest.raises(NotFoundError, match="Unknown resource"):
|
|
await manager.get_resource("greet://world")
|
|
|
|
|
|
class TestResourceErrorHandling:
|
|
"""Test error handling in the ResourceManager."""
|
|
|
|
async def test_resource_error_passthrough(self):
|
|
"""Test that ResourceErrors are passed through directly."""
|
|
manager = ResourceManager()
|
|
|
|
async def error_resource():
|
|
"""Resource that raises a ResourceError."""
|
|
raise ResourceError("Specific resource error")
|
|
|
|
resource = FunctionResource(
|
|
uri=AnyUrl("error://resource"),
|
|
name="error_resource",
|
|
fn=error_resource,
|
|
)
|
|
manager.add_resource(resource)
|
|
|
|
with pytest.raises(ResourceError, match="Specific resource error"):
|
|
await manager.read_resource("error://resource")
|
|
|
|
async def test_template_resource_error_passthrough(self):
|
|
"""Test that ResourceErrors from template-generated resources are passed through."""
|
|
manager = ResourceManager()
|
|
|
|
def error_template(param: str):
|
|
"""Template that raises a ResourceError."""
|
|
raise ResourceError(f"Template error with param {param}")
|
|
|
|
template = ResourceTemplate.from_function(
|
|
fn=error_template,
|
|
uri_template="error://{param}",
|
|
name="error_template",
|
|
)
|
|
manager.add_template(template)
|
|
|
|
with pytest.raises(ResourceError) as excinfo:
|
|
await manager.read_resource("error://test")
|
|
|
|
# The original error message should be included in the ValueError
|
|
assert "Template error with param test" in str(excinfo.value)
|
|
|
|
async def test_exception_converted_to_resource_error_with_details(self):
|
|
"""Test that other exceptions are converted to ResourceError with details by default."""
|
|
manager = ResourceManager()
|
|
|
|
async def buggy_resource():
|
|
"""Resource that raises a ValueError."""
|
|
raise ValueError("Internal error details")
|
|
|
|
resource = FunctionResource(
|
|
uri=AnyUrl("buggy://resource"),
|
|
name="buggy_resource",
|
|
fn=buggy_resource,
|
|
)
|
|
manager.add_resource(resource)
|
|
|
|
with pytest.raises(ResourceError) as excinfo:
|
|
await manager.read_resource("buggy://resource")
|
|
|
|
# The error message should include the original exception details
|
|
assert "Error reading resource 'buggy://resource'" in str(excinfo.value)
|
|
assert "Internal error details" in str(excinfo.value)
|
|
|
|
async def test_exception_converted_to_masked_resource_error(self):
|
|
"""Test that other exceptions are masked when enabled."""
|
|
manager = ResourceManager(mask_error_details=True)
|
|
|
|
async def buggy_resource():
|
|
"""Resource that raises a ValueError."""
|
|
raise ValueError("Internal error details")
|
|
|
|
resource = FunctionResource(
|
|
uri=AnyUrl("buggy://resource"),
|
|
name="buggy_resource",
|
|
fn=buggy_resource,
|
|
)
|
|
manager.add_resource(resource)
|
|
|
|
with pytest.raises(ResourceError) as excinfo:
|
|
await manager.read_resource("buggy://resource")
|
|
|
|
# The error message should not include the original exception details
|
|
assert "Error reading resource 'buggy://resource'" in str(excinfo.value)
|
|
assert "Internal error details" not in str(excinfo.value)
|