Skip to content

Commit fe5e09e

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

1 file changed

Lines changed: 300 additions & 0 deletions

File tree

‎tests/agent/tools/test_extraction_tool.py‎

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

0 commit comments

Comments
 (0)