-
Notifications
You must be signed in to change notification settings - Fork 175
Expand file tree
/
Copy pathtest_tools.py
More file actions
562 lines (488 loc) · 23 KB
/
test_tools.py
File metadata and controls
562 lines (488 loc) · 23 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
"""Tests for all tools.
This module contains simple tests for all tools:
- Initialization
- Config reading (if needed)
"""
import sys
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from sgr_agent_core.models import AgentContext, SourceData
from sgr_agent_core.tools import (
AdaptPlanTool,
AnswerTool,
ClarificationTool,
CreateReportTool,
ExtractPageContentTool,
FinalAnswerTool,
GeneratePlanTool,
ReasoningTool,
RunCommandTool,
WebSearchTool,
)
class TestToolsInitialization:
"""Test that all tools can be initialized."""
def test_clarification_tool_initialization(self):
"""Test ClarificationTool initialization."""
tool = ClarificationTool(
reasoning="Test",
unclear_terms=["term1"],
assumptions=["assumption1", "assumption2"],
questions=["Question 1?", "Question 2?", "Question 3?"],
)
assert tool.tool_name == "clarificationtool"
assert tool.reasoning == "Test"
def test_generate_plan_tool_initialization(self):
"""Test GeneratePlanTool initialization."""
tool = GeneratePlanTool(
reasoning="Test",
research_goal="Test goal",
planned_steps=["Step 1", "Step 2", "Step 3"],
search_strategies=["Strategy 1", "Strategy 2"],
)
assert tool.tool_name == "generateplantool"
assert len(tool.planned_steps) == 3
def test_adapt_plan_tool_initialization(self):
"""Test AdaptPlanTool initialization."""
tool = AdaptPlanTool(
reasoning="Test",
original_goal="Original goal",
new_goal="New goal",
plan_changes=["Change 1"],
next_steps=["Step 1", "Step 2"],
)
assert tool.tool_name == "adaptplantool"
assert len(tool.next_steps) == 2
def test_final_answer_tool_initialization(self):
"""Test FinalAnswerTool initialization."""
from sgr_agent_core.models import AgentStatesEnum
tool = FinalAnswerTool(
reasoning="Test",
completed_steps=["Step 1"],
answer="Answer",
status=AgentStatesEnum.COMPLETED,
)
assert tool.tool_name == "finalanswertool"
assert tool.answer == "Answer"
def test_reasoning_tool_initialization(self):
"""Test ReasoningTool initialization."""
tool = ReasoningTool(
reasoning_steps=["Step 1", "Step 2"],
current_situation="Test",
plan_status="Test",
enough_data=False,
remaining_steps=["Next"],
task_completed=False,
)
assert tool.tool_name == "reasoningtool"
assert len(tool.reasoning_steps) == 2
def test_web_search_tool_initialization(self):
"""Test WebSearchTool initialization."""
tool = WebSearchTool(
reasoning="Test",
query="test query",
)
assert tool.tool_name == "websearchtool"
assert tool.query == "test query"
def test_extract_page_content_tool_initialization(self):
"""Test ExtractPageContentTool initialization."""
tool = ExtractPageContentTool(
reasoning="Test",
urls=["https://example.com"],
)
assert tool.tool_name == "extractpagecontenttool"
assert len(tool.urls) == 1
def test_create_report_tool_initialization(self):
"""Test CreateReportTool initialization."""
tool = CreateReportTool(
reasoning="Test",
title="Test Report",
user_request_language_reference="Test",
content="Test content",
confidence="high",
)
assert tool.tool_name == "createreporttool"
assert tool.title == "Test Report"
def test_answer_tool_initialization(self):
"""Test AnswerTool initialization."""
tool = AnswerTool(
reasoning="Sharing progress",
intermediate_result="Found 3 relevant sources so far.",
continue_research=True,
)
assert tool.tool_name == "answertool"
assert tool.reasoning == "Sharing progress"
assert tool.intermediate_result == "Found 3 relevant sources so far."
assert tool.continue_research is True
class TestAnswerToolExecution:
"""Tests for AnswerTool execution."""
@pytest.mark.asyncio
async def test_answer_tool_returns_intermediate_result(self):
"""Test AnswerTool __call__ returns intermediate_result."""
tool = AnswerTool(
reasoning="Progress update",
intermediate_result="Partial findings: X and Y.",
)
result = await tool(MagicMock(), MagicMock())
assert result == "Partial findings: X and Y."
class TestToolsConfigReading:
"""Test that tools that need config can read it correctly."""
def test_web_search_tool_reads_config(self):
"""Test WebSearchTool reads search config for max_results."""
tool = WebSearchTool(
reasoning="Test",
query="test query",
max_results=5,
)
# Tool should use provided max_results
assert tool.query == "test query"
assert tool.max_results == 5
def test_extract_page_content_tool_reads_config(self):
"""Test ExtractPageContentTool reads search config."""
tool = ExtractPageContentTool(
reasoning="Test",
urls=["https://example.com"],
)
# Tool should be initialized without errors
assert len(tool.urls) == 1
def test_create_report_tool_reads_config(self):
"""Test CreateReportTool reads execution config."""
tool = CreateReportTool(
reasoning="Test",
title="Test Report",
user_request_language_reference="Test",
content="Test content",
confidence="high",
)
# Tool should be initialized without errors
assert tool.title == "Test Report"
class TestSearchToolsKwargs:
"""Test that search tools use kwargs (tool config) for their search
settings."""
@pytest.mark.asyncio
async def test_web_search_tool_uses_kwargs_max_results(self):
"""WebSearchTool uses max_results from kwargs when provided."""
tool = WebSearchTool(reasoning="r", query="test", max_results=5)
context = AgentContext()
config = MagicMock()
mock_handler = AsyncMock(return_value=[])
with patch.dict("sgr_agent_core.tools.web_search_tool._ENGINE_HANDLERS", {"tavily": mock_handler}):
await tool(context, config, api_key="k", max_results=3)
assert mock_handler.call_args.kwargs["max_results"] == 3
@pytest.mark.asyncio
async def test_web_search_tool_default_max_results(self):
"""WebSearchTool uses default max_results when not overridden in
kwargs."""
tool = WebSearchTool(reasoning="r", query="test", max_results=5)
context = AgentContext()
config = MagicMock()
mock_handler = AsyncMock(return_value=[])
with patch.dict("sgr_agent_core.tools.web_search_tool._ENGINE_HANDLERS", {"tavily": mock_handler}):
await tool(context, config, api_key="k")
assert mock_handler.call_args.kwargs["max_results"] == 5
@pytest.mark.asyncio
async def test_web_search_tool_with_offset(self):
"""WebSearchTool passes offset to provider which handles it
internally."""
tool = WebSearchTool(reasoning="r", query="test", max_results=3, offset=2)
context = AgentContext()
config = MagicMock()
mock_sources = [
SourceData(number=i, url=f"https://example.com/{i}", title=f"Result {i}", snippet=f"Snippet {i}")
for i in range(2, 5)
]
mock_handler = AsyncMock(return_value=mock_sources)
with patch.dict("sgr_agent_core.tools.web_search_tool._ENGINE_HANDLERS", {"tavily": mock_handler}):
result = await tool(context, config, api_key="k")
assert len(context.searches) == 1
assert len(context.searches[0].citations) == 3
assert "Result 2" in result
@pytest.mark.asyncio
async def test_web_search_tool_offset_default_zero(self):
"""WebSearchTool without offset passes offset=0 to provider."""
tool = WebSearchTool(reasoning="r", query="test", max_results=3)
assert tool.offset == 0
context = AgentContext()
config = MagicMock()
mock_sources = [
SourceData(number=i, url=f"https://example.com/{i}", title=f"Result {i}", snippet=f"Snippet {i}")
for i in range(3)
]
mock_handler = AsyncMock(return_value=mock_sources)
with patch.dict("sgr_agent_core.tools.web_search_tool._ENGINE_HANDLERS", {"tavily": mock_handler}):
await tool(context, config, api_key="k")
assert len(context.searches[0].citations) == 3
@pytest.mark.asyncio
async def test_web_search_tool_offset_exceeds_results(self):
"""WebSearchTool with offset exceeding available results returns empty
list gracefully."""
tool = WebSearchTool(reasoning="r", query="test", max_results=3, offset=10)
context = AgentContext()
config = MagicMock()
mock_handler = AsyncMock(return_value=[])
with patch.dict("sgr_agent_core.tools.web_search_tool._ENGINE_HANDLERS", {"tavily": mock_handler}):
result = await tool(context, config, api_key="k")
assert len(context.searches[0].citations) == 0
assert "Search Query: test" in result
@pytest.mark.asyncio
async def test_extract_page_content_tool_uses_content_limit_from_kwargs(self):
"""ExtractPageContentTool uses content_limit from kwargs."""
tool = ExtractPageContentTool(reasoning="r", urls=["https://example.com"])
context = AgentContext()
config = MagicMock()
with patch.object(ExtractPageContentTool, "_extract", new_callable=AsyncMock, return_value=[]) as mock_extract:
await tool(context, config, tavily_api_key="k", content_limit=500)
# search_config is passed as first positional arg
assert mock_extract.call_args[0][0].content_limit == 500
class TestRunCommandTool:
"""Test suite for RunCommandTool."""
def test_run_command_tool_initialization(self):
"""RunCommandTool initializes with reasoning and command."""
tool = RunCommandTool(reasoning="List files", command="ls -la")
assert tool.tool_name == "runcommandtool"
assert tool.reasoning == "List files"
assert tool.command == "ls -la"
def test_run_command_tool_default_mode_is_safe(self):
"""RunCommandTool default mode is safe when mode not passed in
kwargs."""
from sgr_agent_core.tools.run_command_tool import RunCommandToolConfig
cfg = RunCommandToolConfig()
assert cfg.mode == "safe"
@pytest.mark.asyncio
async def test_run_command_tool_without_mode_uses_safe_path(self):
"""RunCommandTool without explicit mode uses safe (bwrap) path."""
from sgr_agent_core.models import AgentContext
with tempfile.TemporaryDirectory() as tmpdir:
tool = RunCommandTool(reasoning="Test", command="echo hi")
context = AgentContext()
config = MagicMock()
with (
patch("sgr_agent_core.tools.run_command_tool.shutil.which", return_value="/usr/bin/bwrap"),
patch(
"sgr_agent_core.tools.run_command_tool.asyncio.create_subprocess_exec",
new_callable=AsyncMock,
) as mock_exec,
):
proc = AsyncMock()
proc.communicate = AsyncMock(return_value=(b"hi\n", b""))
proc.returncode = 0
proc.kill = MagicMock()
mock_exec.return_value = proc
result = await tool(context, config, workspace_path=tmpdir)
assert "hi" in result
mock_exec.assert_called_once()
call_args = mock_exec.call_args[0]
assert "bwrap" in str(call_args[0])
@pytest.mark.asyncio
async def test_run_command_tool_unsafe_mode_runs_subprocess(self):
"""RunCommandTool in unsafe mode runs command via subprocess and
returns output."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="echo hello")
context = AgentContext()
config = MagicMock()
result = await tool(context, config, mode="unsafe")
assert "hello" in result
assert "return_code" in result.lower() or "0" in result
@pytest.mark.asyncio
async def test_run_command_tool_uses_workspace_path_as_cwd(self):
"""RunCommandTool with workspace_path runs subprocess with cwd set to
workspace_path (print cwd via same Python as test runner)."""
from sgr_agent_core.models import AgentContext
tmp = Path(__file__).resolve().parent
command = f'"{sys.executable}" -c "import os; print(os.getcwd())"'
tool = RunCommandTool(reasoning="Test", command=command)
context = AgentContext()
config = MagicMock()
result = await tool(context, config, mode="unsafe", workspace_path=str(tmp))
print(result)
print(tmp.name)
assert tmp.name in result or str(tmp) in result
@pytest.mark.asyncio
async def test_run_command_tool_safe_mode_bwrap_not_found_returns_error(self):
"""RunCommandTool in safe mode when bwrap is not installed returns
error with install link."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="echo hi")
context = AgentContext()
config = MagicMock()
with patch("sgr_agent_core.tools.run_command_tool.shutil.which", return_value=None):
result = await tool(context, config, mode="safe")
assert "error" in result.lower()
assert "bwrap" in result.lower()
assert "github.com" in result or "containers/bubblewrap" in result or "installation" in result.lower()
@pytest.mark.asyncio
async def test_run_command_tool_safe_mode_with_bwrap_runs_command(self):
"""RunCommandTool in safe mode with bwrap available runs command via
bwrap."""
from sgr_agent_core.models import AgentContext
with tempfile.TemporaryDirectory() as tmpdir:
tool = RunCommandTool(reasoning="Test", command="echo hi")
context = AgentContext()
config = MagicMock()
with (
patch("sgr_agent_core.tools.run_command_tool.shutil.which", return_value="/usr/bin/bwrap"),
patch(
"sgr_agent_core.tools.run_command_tool.asyncio.create_subprocess_exec",
new_callable=AsyncMock,
) as mock_exec,
):
proc = AsyncMock()
proc.communicate = AsyncMock(return_value=(b"hi\n", b""))
proc.returncode = 0
proc.kill = MagicMock()
mock_exec.return_value = proc
result = await tool(context, config, mode="safe", workspace_path=tmpdir)
assert "hi" in result
assert "return_code" in result.lower()
mock_exec.assert_called_once()
call_args = mock_exec.call_args[0]
assert call_args[0] == "/usr/bin/bwrap" or "bwrap" in str(call_args)
@pytest.mark.asyncio
async def test_run_command_tool_uses_timeout_from_kwargs(self):
"""RunCommandTool uses timeout_seconds from kwargs."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="echo ok")
context = AgentContext()
config = MagicMock()
with patch("sgr_agent_core.tools.run_command_tool.asyncio.create_subprocess_shell") as mock_create:
proc = AsyncMock()
proc.communicate = AsyncMock(return_value=(b"ok\n", b""))
proc.returncode = 0
proc.kill = MagicMock()
mock_create.return_value = proc
await tool(context, config, mode="unsafe", timeout_seconds=30)
call_kwargs = mock_create.call_args[1]
assert call_kwargs.get("cwd") is None or "cwd" in call_kwargs
# timeout is applied in wait_for(communicate(), timeout=...)
mock_create.assert_called_once()
@pytest.mark.asyncio
async def test_run_command_tool_path_escape_rejected_when_workspace_path_set(self):
"""RunCommandTool rejects command that escapes workspace_path (e.g.
../../../)."""
from sgr_agent_core.models import AgentContext
tmp = Path(__file__).resolve().parent
tool = RunCommandTool(reasoning="Test", command="cat ../../../../etc/passwd")
context = AgentContext()
config = MagicMock()
result = await tool(context, config, mode="unsafe", workspace_path=str(tmp))
assert "error" in result.lower() or "not allowed" in result.lower() or "outside" in result.lower()
@pytest.mark.asyncio
async def test_run_command_tool_include_allows_command(self):
"""RunCommandTool with include allows only listed commands."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="echo hello")
context = AgentContext()
config = MagicMock()
# echo should be allowed if in include
result = await tool(context, config, mode="unsafe", include_paths=["echo", "/bin/echo"])
assert "hello" in result
@pytest.mark.asyncio
async def test_run_command_tool_include_rejects_not_listed_command(self):
"""RunCommandTool with include rejects commands not in the list."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="rm /tmp/file")
context = AgentContext()
config = MagicMock()
result = await tool(context, config, mode="unsafe", include_paths=["echo", "ls"])
assert "error" in result.lower()
assert "not in include" in result.lower() or "not allowed" in result.lower()
@pytest.mark.asyncio
async def test_run_command_tool_exclude_rejects_command(self):
"""RunCommandTool with exclude rejects excluded commands."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="rm /tmp/file")
context = AgentContext()
config = MagicMock()
result = await tool(context, config, mode="unsafe", exclude_paths=["rm", "/usr/bin/rm"])
assert "error" in result.lower()
assert "excluded" in result.lower()
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_run_command_tool_include_paths_priority_over_exclude_paths(self):
"""RunCommandTool: include_paths has priority over exclude_paths (same path in both is allowed)."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="rm /tmp/file")
context = AgentContext()
config = MagicMock()
# rm is in both include_paths and exclude_paths -> allowed (include_paths wins)
result = await tool(context, config, mode="unsafe", include_paths=["rm", "ls"], exclude_paths=["rm"])
assert "excluded" not in result.lower()
@pytest.mark.asyncio
async def test_run_command_tool_exclude_paths_rejects_when_not_in_include_paths(self):
"""RunCommandTool: command only in exclude_paths is rejected."""
from sgr_agent_core.models import AgentContext
tool = RunCommandTool(reasoning="Test", command="rm /tmp/file")
context = AgentContext()
config = MagicMock()
result = await tool(context, config, mode="unsafe", include_paths=["ls", "cat"], exclude_paths=["rm"])
assert "error" in result.lower()
assert "excluded" in result.lower()
@pytest.mark.asyncio
async def test_run_command_tool_safe_mode_with_include_uses_overlayfs_manager(self):
"""RunCommandTool in safe mode with include uses OverlayFSManager
mounts."""
from sgr_agent_core.models import AgentContext
with tempfile.TemporaryDirectory() as tmpdir:
tool = RunCommandTool(reasoning="Test", command="echo hi")
context = AgentContext()
config = MagicMock()
with (
patch("sgr_agent_core.tools.run_command_tool.shutil.which", return_value="/usr/bin/bwrap"),
patch(
"sgr_agent_core.tools.run_command_tool.asyncio.create_subprocess_exec",
new_callable=AsyncMock,
) as mock_exec,
patch(
"sgr_agent_core.services.overlayfs_manager.OverlayFSManager.get_overlay_mounts",
return_value={"/usr/bin": "/tmp/merged_usr_bin"},
) as mock_get_mounts,
):
proc = AsyncMock()
proc.communicate = AsyncMock(return_value=(b"hi\n", b""))
proc.returncode = 0
proc.kill = MagicMock()
mock_exec.return_value = proc
# Include only echo (should be in /usr/bin or /bin)
result = await tool(context, config, mode="safe", workspace_path=tmpdir, include_paths=["echo"])
assert "hi" in result
mock_exec.assert_called_once()
mock_get_mounts.assert_called_once()
@pytest.mark.asyncio
async def test_run_command_tool_safe_mode_with_exclude_uses_overlayfs_manager(self):
"""RunCommandTool in safe mode with exclude uses OverlayFSManager
mounts."""
from sgr_agent_core.models import AgentContext
with tempfile.TemporaryDirectory() as tmpdir:
tool = RunCommandTool(reasoning="Test", command="ls")
context = AgentContext()
config = MagicMock()
with (
patch("sgr_agent_core.tools.run_command_tool.shutil.which", return_value="/usr/bin/bwrap"),
patch(
"sgr_agent_core.tools.run_command_tool.asyncio.create_subprocess_exec",
new_callable=AsyncMock,
) as mock_exec,
patch(
"sgr_agent_core.services.overlayfs_manager.OverlayFSManager.get_overlay_mounts",
return_value={"/usr/bin": "/tmp/merged_usr_bin"},
) as mock_get_mounts,
patch("sgr_agent_core.tools.run_command_tool._check_allowed", return_value=None),
):
proc = AsyncMock()
proc.communicate = AsyncMock(return_value=(b"ls\n", b""))
proc.returncode = 0
proc.kill = MagicMock()
mock_exec.return_value = proc
# Include ls but exclude rm (both in /usr/bin)
result = await tool(
context,
config,
mode="safe",
workspace_path=tmpdir,
include_paths=["ls"],
exclude_paths=["rm"],
)
assert "ls" in result
mock_get_mounts.assert_called_once()