123 lines
3.9 KiB
Python
123 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
|
|
class TestDocumentHostTools(unittest.TestCase):
|
|
def test_document_search_truncates_content_without_requiring_key_arg(self):
|
|
from tools.documents import DocumentSearchTool
|
|
|
|
docs = [
|
|
{
|
|
"file_name": "paper.md",
|
|
"kb_name": "mu_1",
|
|
"character_count": 50000,
|
|
"md_content": "A" * 200,
|
|
}
|
|
]
|
|
with patch("tools.documents.doc_client.search", return_value=docs) as search:
|
|
out = DocumentSearchTool().execute(
|
|
query="cement hydration",
|
|
max_documents=3,
|
|
content_chars_per_doc=20,
|
|
)
|
|
|
|
search.assert_called_once_with(
|
|
query="cement hydration",
|
|
kb_names=None,
|
|
classification_ids=None,
|
|
max_documents=3,
|
|
)
|
|
self.assertIn("paper.md", out)
|
|
self.assertIn("mu_1", out)
|
|
self.assertIn("A" * 20, out)
|
|
self.assertIn("truncated", out)
|
|
|
|
def test_document_download_uses_constructor_working_dir(self):
|
|
from tools.documents import DocumentDownloadTool
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
working_dir = Path(tmp) / "task"
|
|
working_dir.mkdir()
|
|
with patch(
|
|
"tools.documents.doc_client.download",
|
|
return_value="documents/paper.pdf",
|
|
) as download:
|
|
tool = DocumentDownloadTool(
|
|
working_dir=working_dir,
|
|
base_dir=working_dir,
|
|
user_root=Path(tmp),
|
|
)
|
|
out = tool.execute(file_name="paper.pdf", kb_name="mu_1")
|
|
|
|
download.assert_called_once_with(
|
|
file_name="paper.pdf",
|
|
kb_name="mu_1",
|
|
working_dir=str(working_dir),
|
|
preview=False,
|
|
)
|
|
self.assertIn("saved: task/documents/paper.pdf", out)
|
|
|
|
|
|
class TestMaterialsProjectHostTools(unittest.TestCase):
|
|
def test_mp_search_summary_uses_host_key_and_returns_json(self):
|
|
from tools.materials_project import MaterialsProjectSearchSummaryTool
|
|
|
|
class FakeDoc:
|
|
material_id = "mp-1"
|
|
formula_pretty = "Ca3SiO5"
|
|
energy_above_hull = 0.0123
|
|
|
|
captured = {}
|
|
|
|
class FakeSummary:
|
|
def search(self, **kwargs):
|
|
captured.update(kwargs)
|
|
return [FakeDoc()]
|
|
|
|
class FakeMaterials:
|
|
def __init__(self):
|
|
self.summary = FakeSummary()
|
|
|
|
class FakeMPRester:
|
|
def __init__(self, api_key):
|
|
self.api_key = api_key
|
|
self.materials = FakeMaterials()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
with patch.dict("os.environ", {"MP_API_KEY": "host-secret"}, clear=False), patch(
|
|
"tools.materials_project.MPRester",
|
|
FakeMPRester,
|
|
):
|
|
out = MaterialsProjectSearchSummaryTool().execute(
|
|
formula="Ca3SiO5",
|
|
fields=["material_id", "formula_pretty", "energy_above_hull"],
|
|
limit=2,
|
|
)
|
|
|
|
data = json.loads(out)
|
|
self.assertEqual(data[0]["material_id"], "mp-1")
|
|
self.assertEqual(data[0]["formula_pretty"], "Ca3SiO5")
|
|
self.assertEqual(data[0]["energy_above_hull"], 0.0123)
|
|
self.assertNotIn("host-secret", out)
|
|
# Server-side limiting: a single page of `limit` records, never the
|
|
# default chunk_size=1000 full-database pagination that gets the IP banned.
|
|
self.assertEqual(captured["num_chunks"], 1)
|
|
self.assertEqual(captured["chunk_size"], 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|