From 29b31573bc696fe26af01976f9053ce097f481a7 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sun, 13 Apr 2025 21:45:17 -0400
Subject: [PATCH] Fix typing issue

---
 tests/prompts/test_base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/prompts/test_base.py b/tests/prompts/test_base.py
index fb02c93d1..84455d068 100644
--- a/tests/prompts/test_base.py
+++ b/tests/prompts/test_base.py
@@ -57,7 +57,7 @@ class TestRenderPrompt:
 
     @pytest.mark.anyio
     async def test_fn_returns_message(self):
-        async def fn() -> UserMessage:
+        async def fn() -> Message:
             return UserMessage(content="Hello, world!")
 
         prompt = Prompt.from_function(fn)
@@ -67,7 +67,7 @@ class TestRenderPrompt:
 
     @pytest.mark.anyio
     async def test_fn_returns_assistant_message(self):
-        async def fn() -> AssistantMessage:
+        async def fn() -> Message:
             return AssistantMessage(
                 content=TextContent(type="text", text="Hello, world!")
             )
@@ -108,7 +108,7 @@ class TestRenderPrompt:
     async def test_fn_returns_resource_content(self):
         """Test returning a message with resource content."""
 
-        async def fn() -> UserMessage:
+        async def fn() -> Message:
             return UserMessage(
                 content=EmbeddedResource(
                     type="resource",