blob: a6c6af593c98ed813662ab941c3edd320d9a4b82 (
plain)
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
|
from __future__ import annotations
from mcp import Client
from mcp.types import ImageContent
from mcp_sd15.server import create_server
from mcp_sd15.service import ImageGenerationService
from .test_service import FakeBackend
async def test_exposes_exactly_three_tools_with_closed_enums(app_config) -> None:
server = create_server(ImageGenerationService(app_config, FakeBackend()))
async with Client(server) as client:
tools = (await client.list_tools()).tools
assert {tool.name for tool in tools} == {
"generate_image",
"generate_guided_image",
"inpaint_image",
}
guided = next(tool for tool in tools if tool.name == "generate_guided_image")
assert guided.input_schema["properties"]["control_mode"]["enum"] == [
"depth",
"openpose",
"openpose_full",
]
assert "width" not in guided.input_schema["properties"]
assert all("scheduler" not in tool.input_schema["properties"] for tool in tools)
async def test_mcp_call_returns_png_image_content(app_config) -> None:
server = create_server(ImageGenerationService(app_config, FakeBackend()))
async with Client(server) as client:
result = await client.call_tool("generate_image", {"positive_prompt": "a cat"})
assert result.is_error is False
assert len(result.content) == 1
assert isinstance(result.content[0], ImageContent)
assert result.content[0].mime_type == "image/png"
|
