Improve coverage
This commit is contained in:
parent
eb4e1721ce
commit
8e555ab76f
2 changed files with 50 additions and 0 deletions
|
|
@ -159,3 +159,36 @@ class TestExecuteCodeTool:
|
||||||
ctx, code="docs = await list_documents()\nprint(len(docs))"
|
ctx, code="docs = await list_documents()\nprint(len(docs))"
|
||||||
)
|
)
|
||||||
assert "1" in result
|
assert "1" in result
|
||||||
|
|
||||||
|
async def test_execute_code_accumulates_search_results(self, rag_db):
|
||||||
|
from haiku.rag.skills.analysis import create_skill
|
||||||
|
|
||||||
|
skill = create_skill(db_path=rag_db)
|
||||||
|
execute_code = _get_tool(skill, "execute_code")
|
||||||
|
state = AnalysisState()
|
||||||
|
ctx = _make_ctx(state)
|
||||||
|
await execute_code(
|
||||||
|
ctx, code="results = await search('intelligence')\nprint(len(results))"
|
||||||
|
)
|
||||||
|
assert "_sandbox" in state.searches
|
||||||
|
assert len(state.searches["_sandbox"]) > 0
|
||||||
|
|
||||||
|
async def test_execute_code_vfs_write_denied(self, rag_db):
|
||||||
|
from haiku.rag.skills.analysis import create_skill
|
||||||
|
|
||||||
|
skill = create_skill(db_path=rag_db)
|
||||||
|
execute_code = _get_tool(skill, "execute_code")
|
||||||
|
state = AnalysisState()
|
||||||
|
ctx = _make_ctx(state)
|
||||||
|
result = await execute_code(
|
||||||
|
ctx,
|
||||||
|
code=(
|
||||||
|
"from pathlib import Path\n"
|
||||||
|
"import json\n"
|
||||||
|
"dirs = list(Path('/documents').iterdir())\n"
|
||||||
|
"p = dirs[0] / 'content.txt'\n"
|
||||||
|
"p.write_text('hacked')"
|
||||||
|
),
|
||||||
|
)
|
||||||
|
assert "Error" in result
|
||||||
|
assert "read-only" in result
|
||||||
|
|
|
||||||
|
|
@ -200,6 +200,23 @@ class TestSearchTool:
|
||||||
result = await search(ctx, query="artificial intelligence")
|
result = await search(ctx, query="artificial intelligence")
|
||||||
assert isinstance(result, str)
|
assert isinstance(result, str)
|
||||||
|
|
||||||
|
async def test_search_rate_limited(self, rag_db):
|
||||||
|
from haiku.rag.skills.rag import RAGState, create_skill
|
||||||
|
|
||||||
|
config = AppConfig()
|
||||||
|
config.qa.max_searches = 2
|
||||||
|
skill = create_skill(db_path=rag_db, config=config)
|
||||||
|
search = _get_tool(skill, "search")
|
||||||
|
state = RAGState()
|
||||||
|
ctx = _make_ctx(state)
|
||||||
|
ctx.run_id = "test-run"
|
||||||
|
|
||||||
|
await search(ctx, query="first")
|
||||||
|
await search(ctx, query="second")
|
||||||
|
result = await search(ctx, query="third")
|
||||||
|
assert "Search limit reached" in result
|
||||||
|
assert len(state.searches) == 2
|
||||||
|
|
||||||
|
|
||||||
class TestListDocumentsTool:
|
class TestListDocumentsTool:
|
||||||
async def test_list_documents_returns_results(self, rag_db):
|
async def test_list_documents_returns_results(self, rag_db):
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue