Skip to content

Commit cdca545

Browse files
committed
add test for requiring context
1 parent 7f1ba04 commit cdca545

File tree

1 file changed

+47
-0
lines changed

1 file changed

+47
-0
lines changed

tests/server/test_completion_with_context.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,3 +131,50 @@ async def handle_completion(
131131
context_arguments={"database": "products_db"},
132132
)
133133
assert table_result2.completion.values == ["products", "categories", "inventory"]
134+
135+
136+
@pytest.mark.anyio
137+
async def test_completion_error_on_missing_context():
138+
"""Test that server can raise error when required context is missing."""
139+
server = Server("test-server")
140+
141+
@server.completion()
142+
async def handle_completion(
143+
ref: PromptReference | ResourceTemplateReference,
144+
argument: CompletionArgument,
145+
context: CompletionContext | None,
146+
) -> Completion | None:
147+
if isinstance(ref, ResourceTemplateReference):
148+
if ref.uri == "db://{database}/{table}":
149+
if argument.name == "table":
150+
# Check if database context is provided
151+
if not context or not context.arguments or "database" not in context.arguments:
152+
# Raise an error instead of returning error as completion
153+
raise ValueError("Please select a database first to see available tables")
154+
# Normal completion if context is provided
155+
db = context.arguments.get("database")
156+
if db == "test_db":
157+
return Completion(values=["users", "orders", "products"], total=3, hasMore=False)
158+
159+
return Completion(values=[], total=0, hasMore=False)
160+
161+
async with create_connected_server_and_client_session(server) as client:
162+
# Try to complete table without database context - should raise error
163+
with pytest.raises(Exception) as exc_info:
164+
await client.complete(
165+
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
166+
argument={"name": "table", "value": ""},
167+
)
168+
169+
# Verify error message
170+
assert "Please select a database first" in str(exc_info.value)
171+
172+
# Now complete with proper context - should work normally
173+
result_with_context = await client.complete(
174+
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
175+
argument={"name": "table", "value": ""},
176+
context_arguments={"database": "test_db"},
177+
)
178+
179+
# Should get normal completions
180+
assert result_with_context.completion.values == ["users", "orders", "products"]

0 commit comments

Comments
 (0)