Skip to content

Commit 52dea3e

Browse files
bendruckerbendrucker
committed
Add edge case tests for tool approval extraction coverage
1 parent e885d43 commit 52dea3e

File tree

1 file changed

+151
-0
lines changed

1 file changed

+151
-0
lines changed

‎tests/test_vercel_ai.py‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1792,6 +1792,19 @@ def delete_file(path: str) -> str:
17921792
id='assistant-1',
17931793
role='assistant',
17941794
parts=[
1795+
# Text part - covers continue branch for non-tool parts in _denied_tool_ids
1796+
TextUIPart(text='I will delete the file for you.'),
1797+
# Approved tool - covers branch 100->97 (approved=True) in _denied_tool_ids
1798+
DynamicToolInputAvailablePart(
1799+
tool_name='delete_file',
1800+
tool_call_id='delete_approved',
1801+
input={'path': 'approved.txt'},
1802+
approval=ToolApprovalResponded(
1803+
id='approval-456',
1804+
approved=True,
1805+
),
1806+
),
1807+
# Denied tool - should be processed
17951808
DynamicToolInputAvailablePart(
17961809
tool_name='delete_file',
17971810
tool_call_id='delete_1',
@@ -1832,6 +1845,144 @@ def mock_header_get(key: str) -> str | None:
18321845
assert denied_event['toolCallId'] == 'delete_1'
18331846

18341847

1848+
@pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed')
1849+
async def test_tool_approval_extraction_with_edge_cases():
1850+
"""Test that approval extraction correctly skips non-tool parts and non-responded approvals."""
1851+
from unittest.mock import AsyncMock
1852+
1853+
from starlette.requests import Request
1854+
1855+
from pydantic_ai.tools import DeferredToolRequests
1856+
from pydantic_ai.ui.vercel_ai.request_types import (
1857+
DynamicToolInputAvailablePart,
1858+
ToolApprovalRequested,
1859+
ToolApprovalResponded,
1860+
)
1861+
1862+
agent = Agent(TestModel(), output_type=[str, DeferredToolRequests])
1863+
1864+
@agent.tool_plain(requires_approval=True)
1865+
def some_tool(x: str) -> str:
1866+
return x # pragma: no cover
1867+
1868+
# Request with edge cases: TextUIPart (non-tool part) and ToolApprovalRequested (non-responded)
1869+
request = SubmitMessage(
1870+
id='foo',
1871+
messages=[
1872+
UIMessage(
1873+
id='user-1',
1874+
role='user',
1875+
parts=[TextUIPart(text='Test')],
1876+
),
1877+
UIMessage(
1878+
id='assistant-1',
1879+
role='assistant',
1880+
parts=[
1881+
# Text part in assistant message - should be skipped
1882+
TextUIPart(text='Here is my response.'),
1883+
# Tool with pending approval (not yet responded) - should be skipped
1884+
DynamicToolInputAvailablePart(
1885+
tool_name='some_tool',
1886+
tool_call_id='pending_tool',
1887+
input={'x': 'pending'},
1888+
approval=ToolApprovalRequested(id='pending-approval'),
1889+
),
1890+
# Tool with no approval at all - should be skipped
1891+
DynamicToolInputAvailablePart(
1892+
tool_name='some_tool',
1893+
tool_call_id='no_approval_tool',
1894+
input={'x': 'no_approval'},
1895+
approval=None,
1896+
),
1897+
# Tool with responded approval - should be processed
1898+
DynamicToolInputAvailablePart(
1899+
tool_name='some_tool',
1900+
tool_call_id='approved_tool',
1901+
input={'x': 'approved'},
1902+
approval=ToolApprovalResponded(id='approved-id', approved=True),
1903+
),
1904+
],
1905+
),
1906+
],
1907+
)
1908+
1909+
def mock_header_get(key: str) -> str | None:
1910+
return None
1911+
1912+
request_body = request.model_dump_json().encode()
1913+
mock_request = AsyncMock(spec=Request)
1914+
mock_request.body = AsyncMock(return_value=request_body)
1915+
mock_request.headers.get = mock_header_get
1916+
1917+
adapter = await VercelAIAdapter[None, str | DeferredToolRequests].from_request(
1918+
mock_request, agent=agent, tool_approval=True
1919+
)
1920+
1921+
# Verify that only the responded approval was extracted
1922+
assert adapter.deferred_tool_results is not None
1923+
assert adapter.deferred_tool_results.approvals == {'approved_tool': True}
1924+
1925+
1926+
@pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed')
1927+
async def test_tool_approval_no_approvals_extracted():
1928+
"""Test that deferred_tool_results is None when no approvals are responded."""
1929+
from unittest.mock import AsyncMock
1930+
1931+
from starlette.requests import Request
1932+
1933+
from pydantic_ai.tools import DeferredToolRequests
1934+
from pydantic_ai.ui.vercel_ai.request_types import (
1935+
DynamicToolInputAvailablePart,
1936+
ToolApprovalRequested,
1937+
)
1938+
1939+
agent = Agent(TestModel(), output_type=[str, DeferredToolRequests])
1940+
1941+
@agent.tool_plain(requires_approval=True)
1942+
def some_tool(x: str) -> str:
1943+
return x # pragma: no cover
1944+
1945+
# Request with only pending approvals (no ToolApprovalResponded)
1946+
request = SubmitMessage(
1947+
id='foo',
1948+
messages=[
1949+
UIMessage(
1950+
id='user-1',
1951+
role='user',
1952+
parts=[TextUIPart(text='Test')],
1953+
),
1954+
UIMessage(
1955+
id='assistant-1',
1956+
role='assistant',
1957+
parts=[
1958+
# Only pending approvals - should not be extracted
1959+
DynamicToolInputAvailablePart(
1960+
tool_name='some_tool',
1961+
tool_call_id='pending_tool',
1962+
input={'x': 'pending'},
1963+
approval=ToolApprovalRequested(id='pending-approval'),
1964+
),
1965+
],
1966+
),
1967+
],
1968+
)
1969+
1970+
def mock_header_get(key: str) -> str | None:
1971+
return None
1972+
1973+
request_body = request.model_dump_json().encode()
1974+
mock_request = AsyncMock(spec=Request)
1975+
mock_request.body = AsyncMock(return_value=request_body)
1976+
mock_request.headers.get = mock_header_get
1977+
1978+
adapter = await VercelAIAdapter[None, str | DeferredToolRequests].from_request(
1979+
mock_request, agent=agent, tool_approval=True
1980+
)
1981+
1982+
# No approvals extracted, so deferred_tool_results should be None
1983+
assert adapter.deferred_tool_results is None
1984+
1985+
18351986
@pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed')
18361987
async def test_adapter_dispatch_request():
18371988
agent = Agent(model=TestModel())

0 commit comments

Comments
 (0)