Skip to content

Commit 2b9f41e

Browse files
radu-mocanuradu-mocanu
committed
chore: add extraction tool tests
1 parent 1d3b99c commit 2b9f41e

1 file changed

Lines changed: 298 additions & 0 deletions

File tree

‎tests/agent/tools/test_extraction_tool.py‎

Lines changed: 298 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
"""Tests for extraction_tool.py metadata and functionality."""
2+
3+
from unittest.mock import AsyncMock, MagicMock, patch
4+
5+
import pytest
6+
from uipath.agent.models.agent import (
7+
AgentIxpExtractionResourceConfig,
8+
AgentIxpExtractionToolProperties,
9+
)
10+
from uipath.platform.attachments import Attachment
11+
from uipath.platform.documents import ExtractionResponseIXP
12+
13+
from uipath_langchain.agent.tools.extraction_tool import create_ixp_extraction_tool
14+
15+
16+
class TestExtractionToolMetadata:
17+
"""Test that extraction tool has correct metadata for observability."""
18+
19+
@pytest.fixture
20+
def extraction_resource(self):
21+
"""Create a minimal extraction tool resource config."""
22+
return AgentIxpExtractionResourceConfig(
23+
name="test_extraction",
24+
description="Extract data from files",
25+
input_schema={
26+
"type": "object",
27+
"properties": {
28+
"attachment": {
29+
"description": "the file uploaded as attachement",
30+
"$ref": "#/definitions/job-attachment",
31+
}
32+
},
33+
"required": ["attachment"],
34+
"definitions": {
35+
"job-attachment": {
36+
"type": "object",
37+
"required": ["ID"],
38+
"x-uipath-resource-kind": "JobAttachment",
39+
"properties": {
40+
"ID": {
41+
"type": "string",
42+
"description": "Orchestrator attachment key",
43+
},
44+
"FullName": {"type": "string", "description": "File name"},
45+
"MimeType": {
46+
"type": "string",
47+
"description": 'The MIME type of the content, such as "application/json" or "image/png"',
48+
},
49+
"Metadata": {
50+
"type": "object",
51+
"description": "Dictionary<string, string> of metadata",
52+
"additionalProperties": {"type": "string"},
53+
},
54+
},
55+
}
56+
},
57+
},
58+
output_schema={"type": "object", "properties": {}},
59+
properties=AgentIxpExtractionToolProperties(
60+
project_name="TestProject",
61+
version_tag="v1.0",
62+
),
63+
)
64+
65+
def test_extraction_tool_has_correct_name(self, extraction_resource):
66+
"""Test that extraction tool has sanitized name."""
67+
tool = create_ixp_extraction_tool(extraction_resource)
68+
69+
assert tool.name == "test_extraction"
70+
71+
def test_extraction_tool_has_correct_description(self, extraction_resource):
72+
"""Test that extraction tool has correct description."""
73+
tool = create_ixp_extraction_tool(extraction_resource)
74+
75+
assert tool.description == "Extract data from files"
76+
77+
def test_extraction_tool_has_attachment_input_schema(self, extraction_resource):
78+
"""Test that extraction tool uses Attachment as input schema."""
79+
tool = create_ixp_extraction_tool(extraction_resource)
80+
81+
assert tool.args_schema == Attachment
82+
83+
def test_extraction_tool_has_extraction_response_output_type(
84+
self, extraction_resource
85+
):
86+
"""Test that extraction tool has ExtractionResponseIXP as output type."""
87+
tool = create_ixp_extraction_tool(extraction_resource)
88+
89+
assert hasattr(tool, "output_type")
90+
assert tool.output_type == ExtractionResponseIXP
91+
92+
93+
class TestExtractionToolFunctionality:
94+
"""Test the extraction tool function behavior."""
95+
96+
@pytest.fixture
97+
def extraction_resource(self):
98+
"""Create a minimal extraction tool resource config."""
99+
return AgentIxpExtractionResourceConfig(
100+
name="test_extraction",
101+
description="Extract data from files",
102+
input_schema={
103+
"type": "object",
104+
"properties": {
105+
"attachment": {
106+
"description": "the file uploaded as attachment",
107+
"$ref": "#/definitions/job-attachment",
108+
}
109+
},
110+
"required": ["attachment"],
111+
"definitions": {
112+
"job-attachment": {
113+
"type": "object",
114+
"required": ["ID"],
115+
"x-uipath-resource-kind": "JobAttachment",
116+
"properties": {
117+
"ID": {
118+
"type": "string",
119+
"description": "Orchestrator attachment key",
120+
},
121+
"FullName": {"type": "string", "description": "File name"},
122+
"MimeType": {
123+
"type": "string",
124+
"description": "The MIME type of the content",
125+
},
126+
},
127+
}
128+
},
129+
},
130+
output_schema={"type": "object", "properties": {}},
131+
properties=AgentIxpExtractionToolProperties(
132+
project_name="TestProject",
133+
version_tag="v1.0",
134+
),
135+
)
136+
137+
@pytest.mark.asyncio
138+
@patch("uipath_langchain.agent.tools.extraction_tool.UiPath")
139+
@patch("uipath_langchain.agent.tools.extraction_tool.interrupt")
140+
async def test_extraction_tool_downloads_attachment_and_calls_interrupt(
141+
self, mock_interrupt, mock_uipath_class, extraction_resource
142+
):
143+
"""Test that extraction tool downloads attachment and calls interrupt with correct params."""
144+
mock_client = MagicMock()
145+
mock_uipath_class.return_value = mock_client
146+
mock_client.attachments.download_async = AsyncMock(
147+
return_value="/path/to/document.pdf"
148+
)
149+
mock_interrupt.return_value = {"extracted_data": {"field1": "value1"}}
150+
151+
tool = create_ixp_extraction_tool(extraction_resource)
152+
153+
result = await tool.ainvoke(
154+
{
155+
"id": "att-123",
156+
"full_name": "document.pdf",
157+
"mime_type": "application/pdf",
158+
}
159+
)
160+
161+
mock_client.attachments.download_async.assert_called_once_with(
162+
key="att-123", destination_path="document.pdf"
163+
)
164+
165+
assert mock_interrupt.called
166+
interrupt_arg = mock_interrupt.call_args[0][0]
167+
assert interrupt_arg.project_name == "TestProject"
168+
assert interrupt_arg.tag == "v1.0"
169+
assert interrupt_arg.file_path == "/path/to/document.pdf"
170+
171+
assert result == {"extracted_data": {"field1": "value1"}}
172+
173+
@pytest.mark.asyncio
174+
@patch("uipath_langchain.agent.tools.extraction_tool.UiPath")
175+
@patch("uipath_langchain.agent.tools.extraction_tool.interrupt")
176+
async def test_extraction_tool_handles_missing_attachment_id(
177+
self, mock_interrupt, mock_uipath_class, extraction_resource
178+
):
179+
"""Test that extraction tool handles None attachment_id."""
180+
mock_client = MagicMock()
181+
mock_uipath_class.return_value = mock_client
182+
mock_client.attachments.download_async = AsyncMock(
183+
return_value="/path/to/file.pdf"
184+
)
185+
mock_interrupt.return_value = {"extracted_data": {}}
186+
187+
tool = create_ixp_extraction_tool(extraction_resource)
188+
189+
await tool.ainvoke({"full_name": "file.pdf", "mime_type": "application/pdf"})
190+
191+
mock_client.attachments.download_async.assert_called_once_with(
192+
key=None, destination_path="file.pdf"
193+
)
194+
195+
@pytest.mark.asyncio
196+
@patch("uipath_langchain.agent.tools.extraction_tool.UiPath")
197+
@patch("uipath_langchain.agent.tools.extraction_tool.interrupt")
198+
async def test_extraction_tool_with_different_version_tag(
199+
self, mock_interrupt, mock_uipath_class
200+
):
201+
"""Test extraction tool with different version tag."""
202+
extraction_resource = AgentIxpExtractionResourceConfig(
203+
name="test_extraction_v2",
204+
description="Extract data from files v2",
205+
input_schema={"type": "object", "properties": {}},
206+
output_schema={"type": "object", "properties": {}},
207+
properties=AgentIxpExtractionToolProperties(
208+
project_name="TestProjectV2",
209+
version_tag="staging",
210+
),
211+
)
212+
213+
mock_client = MagicMock()
214+
mock_uipath_class.return_value = mock_client
215+
mock_client.attachments.download_async = AsyncMock(
216+
return_value="/path/to/document.pdf"
217+
)
218+
mock_interrupt.return_value = {"extracted_data": {}}
219+
220+
tool = create_ixp_extraction_tool(extraction_resource)
221+
222+
await tool.ainvoke(
223+
{
224+
"id": "att-456",
225+
"full_name": "document.pdf",
226+
"mime_type": "application/pdf",
227+
}
228+
)
229+
230+
interrupt_arg = mock_interrupt.call_args[0][0]
231+
assert interrupt_arg.tag == "staging"
232+
233+
@pytest.mark.asyncio
234+
@patch("uipath_langchain.agent.tools.extraction_tool.UiPath")
235+
async def test_extraction_tool_propagates_download_exception(
236+
self, mock_uipath_class, extraction_resource
237+
):
238+
"""Test that exceptions from attachment download are propagated."""
239+
mock_client = MagicMock()
240+
mock_uipath_class.return_value = mock_client
241+
mock_client.attachments.download_async = AsyncMock(
242+
side_effect=Exception("Download failed")
243+
)
244+
245+
tool = create_ixp_extraction_tool(extraction_resource)
246+
247+
with pytest.raises(Exception) as exc_info:
248+
await tool.ainvoke(
249+
{
250+
"id": "att-789",
251+
"full_name": "file.pdf",
252+
"mime_type": "application/pdf",
253+
}
254+
)
255+
256+
assert "Download failed" in str(exc_info.value)
257+
258+
259+
class TestExtractionToolNameSanitization:
260+
"""Test that extraction tool names are properly sanitized."""
261+
262+
@pytest.mark.asyncio
263+
async def test_extraction_tool_name_with_spaces(self):
264+
"""Test that tool names with spaces are sanitized."""
265+
resource = AgentIxpExtractionResourceConfig(
266+
name="Invoice Extraction Tool",
267+
description="Extract invoices",
268+
input_schema={"type": "object", "properties": {}},
269+
output_schema={"type": "object", "properties": {}},
270+
properties=AgentIxpExtractionToolProperties(
271+
project_name="InvoiceExtraction",
272+
version_tag="v1.0",
273+
),
274+
)
275+
276+
tool = create_ixp_extraction_tool(resource)
277+
278+
assert " " not in tool.name
279+
280+
@pytest.mark.asyncio
281+
async def test_extraction_tool_name_with_special_chars(self):
282+
"""Test that tool names with special characters are sanitized."""
283+
resource = AgentIxpExtractionResourceConfig(
284+
name="invoice-extraction@v1",
285+
description="Extract invoices",
286+
input_schema={"type": "object", "properties": {}},
287+
output_schema={"type": "object", "properties": {}},
288+
properties=AgentIxpExtractionToolProperties(
289+
project_name="InvoiceExtraction",
290+
version_tag="v1.0",
291+
),
292+
)
293+
294+
tool = create_ixp_extraction_tool(resource)
295+
296+
# Tool name should be sanitized
297+
assert tool.name is not None
298+
assert len(tool.name) > 0

0 commit comments

Comments
 (0)