49b9bb6724
Deploy Docs / deploy-docs (push) Failing after 1s
Conformance Tests / client-conformance (push) Failing after 3s
Conformance Tests / server-conformance (push) Failing after 1s
GitHub Actions Security Analysis / zizmor (push) Failing after 1s
CI / checks (push) Failing after 59m20s
CI / all-green (push) Has been cancelled
111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
import pytest
|
|
from mcp_types import TextContent
|
|
|
|
from mcp.server.mcpserver import Context
|
|
from mcp.server.mcpserver.prompts.base import Prompt, UserMessage
|
|
from mcp.server.mcpserver.prompts.manager import PromptManager
|
|
|
|
|
|
class TestPromptManager:
|
|
def test_add_prompt(self):
|
|
"""Test adding a prompt to the manager."""
|
|
|
|
def fn() -> str: # pragma: no cover
|
|
return "Hello, world!"
|
|
|
|
manager = PromptManager()
|
|
prompt = Prompt.from_function(fn)
|
|
added = manager.add_prompt(prompt)
|
|
assert added == prompt
|
|
assert manager.get_prompt("fn") == prompt
|
|
|
|
def test_add_duplicate_prompt(self, caplog: pytest.LogCaptureFixture):
|
|
"""Test adding the same prompt twice."""
|
|
|
|
def fn() -> str: # pragma: no cover
|
|
return "Hello, world!"
|
|
|
|
manager = PromptManager()
|
|
prompt = Prompt.from_function(fn)
|
|
first = manager.add_prompt(prompt)
|
|
second = manager.add_prompt(prompt)
|
|
assert first == second
|
|
assert "Prompt already exists" in caplog.text
|
|
|
|
def test_disable_warn_on_duplicate_prompts(self, caplog: pytest.LogCaptureFixture):
|
|
"""Test disabling warning on duplicate prompts."""
|
|
|
|
def fn() -> str: # pragma: no cover
|
|
return "Hello, world!"
|
|
|
|
manager = PromptManager(warn_on_duplicate_prompts=False)
|
|
prompt = Prompt.from_function(fn)
|
|
first = manager.add_prompt(prompt)
|
|
second = manager.add_prompt(prompt)
|
|
assert first == second
|
|
assert "Prompt already exists" not in caplog.text
|
|
|
|
def test_list_prompts(self):
|
|
"""Test listing all prompts."""
|
|
|
|
def fn1() -> str: # pragma: no cover
|
|
return "Hello, world!"
|
|
|
|
def fn2() -> str: # pragma: no cover
|
|
return "Goodbye, world!"
|
|
|
|
manager = PromptManager()
|
|
prompt1 = Prompt.from_function(fn1)
|
|
prompt2 = Prompt.from_function(fn2)
|
|
manager.add_prompt(prompt1)
|
|
manager.add_prompt(prompt2)
|
|
prompts = manager.list_prompts()
|
|
assert len(prompts) == 2
|
|
assert prompts == [prompt1, prompt2]
|
|
|
|
@pytest.mark.anyio
|
|
async def test_render_prompt(self):
|
|
"""Test rendering a prompt."""
|
|
|
|
def fn() -> str:
|
|
return "Hello, world!"
|
|
|
|
manager = PromptManager()
|
|
prompt = Prompt.from_function(fn)
|
|
manager.add_prompt(prompt)
|
|
messages = await manager.render_prompt("fn", None, Context())
|
|
assert messages == [UserMessage(content=TextContent(type="text", text="Hello, world!"))]
|
|
|
|
@pytest.mark.anyio
|
|
async def test_render_prompt_with_args(self):
|
|
"""Test rendering a prompt with arguments."""
|
|
|
|
def fn(name: str) -> str:
|
|
return f"Hello, {name}!"
|
|
|
|
manager = PromptManager()
|
|
prompt = Prompt.from_function(fn)
|
|
manager.add_prompt(prompt)
|
|
messages = await manager.render_prompt("fn", {"name": "World"}, Context())
|
|
assert messages == [UserMessage(content=TextContent(type="text", text="Hello, World!"))]
|
|
|
|
@pytest.mark.anyio
|
|
async def test_render_unknown_prompt(self):
|
|
"""Test rendering a non-existent prompt."""
|
|
manager = PromptManager()
|
|
with pytest.raises(ValueError, match="Unknown prompt: unknown"):
|
|
await manager.render_prompt("unknown", None, Context())
|
|
|
|
@pytest.mark.anyio
|
|
async def test_render_prompt_with_missing_args(self):
|
|
"""Test rendering a prompt with missing required arguments."""
|
|
|
|
def fn(name: str) -> str: # pragma: no cover
|
|
return f"Hello, {name}!"
|
|
|
|
manager = PromptManager()
|
|
prompt = Prompt.from_function(fn)
|
|
manager.add_prompt(prompt)
|
|
with pytest.raises(ValueError, match="Missing required arguments"):
|
|
await manager.render_prompt("fn", None, Context())
|