From 23f78733148795cc21c45d6d4a74e08f8fc1710d Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Fri, 1 Aug 2025 12:55:05 -0700
Subject: [PATCH] Add unit test for sampling with image messages (#1329)

---
 tests/client/test_sampling.py | 63 ++++++++++++++++++++++++++++++++---
 1 file changed, 59 insertions(+), 4 deletions(-)

diff --git a/tests/client/test_sampling.py b/tests/client/test_sampling.py
index 5b11b0885..4c35c209c 100644
--- a/tests/client/test_sampling.py
+++ b/tests/client/test_sampling.py
@@ -1,10 +1,12 @@
-from typing import cast
+import json
 
 import pytest
 from mcp.types import TextContent
+from pydantic_core import to_json
 
 from fastmcp import Client, Context, FastMCP
 from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
+from fastmcp.utilities.types import Image
 
 
 @pytest.fixture
@@ -14,12 +16,12 @@ def fastmcp_server():
     @mcp.tool
     async def simple_sample(message: str, context: Context) -> str:
         result = await context.sample("Hello, world!")
-        return cast(TextContent, result).text
+        return result.text  # type: ignore[attr-defined]
 
     @mcp.tool
     async def sample_with_system_prompt(message: str, context: Context) -> str:
         result = await context.sample("Hello, world!", system_prompt="You love FastMCP")
-        return cast(TextContent, result).text
+        return result.text  # type: ignore[attr-defined]
 
     @mcp.tool
     async def sample_with_messages(message: str, context: Context) -> str:
@@ -34,7 +36,25 @@ def fastmcp_server():
                 ),
             ]
         )
-        return cast(TextContent, result).text
+        return result.text  # type: ignore[attr-defined]
+
+    @mcp.tool
+    async def sample_with_image(image_bytes: bytes, context: Context) -> str:
+        image = Image(data=image_bytes)
+
+        result = await context.sample(
+            [
+                SamplingMessage(
+                    content=TextContent(type="text", text="What's in this image?"),
+                    role="user",
+                ),
+                SamplingMessage(
+                    content=image.to_image_content(),
+                    role="user",
+                ),
+            ]
+        )
+        return result.text  # type: ignore[attr-defined]
 
     return mcp
 
@@ -80,3 +100,38 @@ async def test_sampling_with_messages(fastmcp_server: FastMCP):
             "sample_with_messages", {"message": "Hello, world!"}
         )
         assert result.data == "I need to think."
+
+
+async def test_sampling_with_image(fastmcp_server: FastMCP):
+    def sampling_handler(
+        messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
+    ) -> str:
+        assert len(messages) == 2
+        return to_json(messages).decode()
+
+    async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
+        image_bytes = b"abc123"
+        result = await client.call_tool(
+            "sample_with_image", {"image_bytes": image_bytes}
+        )
+        assert json.loads(result.data) == [
+            {
+                "role": "user",
+                "content": {
+                    "type": "text",
+                    "text": "What's in this image?",
+                    "annotations": None,
+                    "_meta": None,
+                },
+            },
+            {
+                "role": "user",
+                "content": {
+                    "type": "image",
+                    "data": "YWJjMTIz",
+                    "mimeType": "image/png",
+                    "annotations": None,
+                    "_meta": None,
+                },
+            },
+        ]