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 class FakeSummary: def search(self, **kwargs): self.kwargs = 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) if __name__ == "__main__": unittest.main()