import json import tempfile import unittest from pathlib import Path from obsidian_wiki.file_assets import connect, init_db, upsert_asset from obsidian_wiki.mcp_assets import ( AssetAccessError, AssetTooLargeError, asset_uri_for_file_id, file_id_from_asset_uri, get_file_asset_metadata, read_file_asset_resource, search_file_assets, ) class McpAssetTests(unittest.TestCase): def setUp(self): self.tempdir = tempfile.TemporaryDirectory() self.vault = Path(self.tempdir.name) (self.vault / "_meta").mkdir() self.db_path = self.vault / "_meta" / "file-assets.sqlite" self.registry_path = self.vault / "_meta" / "source-registry.jsonl" init_db(self.db_path) self.drawing_id = "sha256:" + "a" * 64 self.order_id = "sha256:" + "b" * 64 self.complaint_id = "sha256:" + "c" * 64 self._write_asset_file("drawings/a.pdf", b"%PDF allowed drawing") self._write_asset_file("orders/order.xlsx", b"customer price data") self._write_asset_file("quality/photo.png", b"complaint evidence") self._insert_asset( self.drawing_id, archive_path="_sources/archive/drawings/a.pdf", original_name="B3501_11101c_X2.pdf", doc_type="drawing", source_type="pdf", product_lines=["BTB"], business_categories=["products/btb", "drawings/customer"], data_classification="technical-confidential", visibility="internal", summary="B3501 customer drawing.", bjd_part_numbers=["101049-X0XXX0"], series=["0.5-18"], pitch=["0.5"], pin_count=["18"], ) self._insert_asset( self.order_id, archive_path="_sources/archive/orders/order.xlsx", original_name="订单项.xlsx", doc_type="order_data", source_type="spreadsheet", business_categories=["orders/raw-excel"], data_classification="commercial-confidential", visibility="pii", summary="Raw customer order lines.", ) self._insert_asset( self.complaint_id, archive_path="_sources/archive/quality/photo.png", original_name="客户反馈来料标签错误图片.png", doc_type="evidence_media", source_type="image", business_categories=["quality/customer-complaints"], data_classification="quality-confidential", visibility="internal", summary="Raw complaint evidence image.", ) def tearDown(self): self.tempdir.cleanup() def _write_asset_file(self, rel_path, data): path = self.vault / "_sources" / "archive" / rel_path path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(data) def _insert_asset(self, file_id, **fields): record = { "record_type": "source_archive", "source_id": file_id, "project": "bjd", "title": Path(fields["original_name"]).stem, "size_bytes": len((self.vault / fields["archive_path"]).read_bytes()), "related_pages": ["projects/bjd/references/example.md"], "allowed_bots": ["bjd-general-bot"], **fields, } with self.registry_path.open("a", encoding="utf-8") as fh: fh.write(json.dumps(record, ensure_ascii=False) + "\n") asset = { "file_id": file_id, "sha256": file_id.removeprefix("sha256:"), "archive_path": fields["archive_path"], "original_name": fields["original_name"], "file_ext": Path(fields["original_name"]).suffix.lower(), "mime_type": "application/pdf", "source_type": fields["source_type"], "doc_type": fields["doc_type"], "project": "bjd", "title": record["title"], "summary": fields["summary"], "visibility": fields["visibility"], "business_categories": json.dumps(fields.get("business_categories", []), ensure_ascii=False), "product_lines": json.dumps(fields.get("product_lines", []), ensure_ascii=False), "customers": "[]", "competitors": "[]", "standards": "[]", "market_tags": "[]", "bjd_part_numbers": json.dumps(fields.get("bjd_part_numbers", []), ensure_ascii=False), "competitor_part_numbers": json.dumps(fields.get("competitor_part_numbers", []), ensure_ascii=False), "series": json.dumps(fields.get("series", []), ensure_ascii=False), "pitch": json.dumps(fields.get("pitch", []), ensure_ascii=False), "pin_count": json.dumps(fields.get("pin_count", []), ensure_ascii=False), "related_pages": json.dumps(record["related_pages"], ensure_ascii=False), "catalog_paths": "[]", "extracted_text": None, "classification_confidence": 0.9, "classification_evidence": "[]", "needs_review": 0, "size_bytes": record["size_bytes"], "modified_at": None, "archived_at": "2026-07-05T00:00:00Z", } with connect(self.db_path) as conn: upsert_asset(conn, asset, status="active") def _write_profile_file(self): data = { "version": 2, "profiles": { "bjd-general-bot": { "allowed_visibility": ["public", "internal"], "allowed_data_classification": [ "public", "internal", "technical-confidential", ], "allowed_business_domains": ["products", "engineering"], "blocked_doc_types": ["order_data", "evidence_media"], "include_source_files": False, } }, } (self.vault / "_meta" / "bot-profiles.json").write_text(json.dumps(data), encoding="utf-8") def test_search_returns_allowed_engineering_assets_only(self): rows = search_file_assets(self.vault, "B3501 drawing", db_path=self.db_path) self.assertEqual([row["file_id"] for row in rows], [self.drawing_id]) self.assertEqual(rows[0]["resource_uri"], asset_uri_for_file_id(self.drawing_id)) self.assertNotIn("absolute_path", rows[0]) def test_search_exposes_safe_identifier_fields(self): rows = search_file_assets(self.vault, "101049 图纸", db_path=self.db_path) self.assertEqual([row["file_id"] for row in rows], [self.drawing_id]) self.assertEqual(rows[0]["bjd_part_numbers"], ["101049-X0XXX0"]) self.assertEqual(rows[0]["series"], ["0.5-18"]) self.assertEqual(rows[0]["pitch"], ["0.5"]) self.assertEqual(rows[0]["pin_count"], ["18"]) def test_customer_shareable_engineering_files_work_with_source_file_profile(self): self._write_profile_file() shareable_id = "sha256:" + "d" * 64 self._write_asset_file("drawings/customer-shareable.pdf", b"%PDF customer drawing") self._insert_asset( shareable_id, archive_path="_sources/archive/drawings/customer-shareable.pdf", original_name="北极大-0.5-18客户图.pdf", doc_type="drawing", source_type="pdf", product_lines=["FPC"], business_categories=["products/fpc", "drawings/customer-drawings"], data_classification="customer-shareable", visibility="public", summary="FPC 0.5-18 customer-shareable drawing.", bjd_part_numbers=["101049-X0XXX0"], series=["0.5-18"], ) rows = search_file_assets(self.vault, "101049 图纸", db_path=self.db_path, profile_name="bjd-general-bot") self.assertEqual({row["file_id"] for row in rows}, {shareable_id, self.drawing_id}) def test_search_blocks_sensitive_order_and_complaint_sources(self): self.assertEqual(search_file_assets(self.vault, "订单", db_path=self.db_path), []) self.assertEqual(search_file_assets(self.vault, "来料标签错误", db_path=self.db_path), []) def test_asset_uri_round_trips_file_id(self): uri = asset_uri_for_file_id(self.drawing_id) self.assertEqual(uri, "asset://source/sha256/" + "a" * 64) self.assertEqual(file_id_from_asset_uri(uri), self.drawing_id) def test_resource_read_rechecks_acl_and_returns_file_bytes(self): data = read_file_asset_resource( self.vault, asset_uri_for_file_id(self.drawing_id), db_path=self.db_path, ) self.assertEqual(data["name"], "B3501_11101c_X2.pdf") self.assertEqual(data["mime_type"], "application/pdf") self.assertEqual(data["bytes"], b"%PDF allowed drawing") def test_resource_read_rejects_blocked_asset_even_by_direct_uri(self): with self.assertRaises(AssetAccessError): read_file_asset_resource( self.vault, asset_uri_for_file_id(self.order_id), db_path=self.db_path, ) def test_resource_read_enforces_max_size(self): with self.assertRaises(AssetTooLargeError): read_file_asset_resource( self.vault, asset_uri_for_file_id(self.drawing_id), db_path=self.db_path, max_file_bytes=5, ) def test_metadata_uses_safe_fields(self): metadata = get_file_asset_metadata(self.vault, self.drawing_id, db_path=self.db_path) self.assertEqual(metadata["file_id"], self.drawing_id) self.assertEqual(metadata["doc_type"], "drawing") self.assertEqual(metadata["resource_uri"], asset_uri_for_file_id(self.drawing_id)) self.assertNotIn("absolute_path", metadata) def test_registry_allowed_bots_are_enforced_without_profile_file(self): with self.assertRaises(AssetAccessError): get_file_asset_metadata( self.vault, self.drawing_id, db_path=self.db_path, profile_name="public-product-bot", ) if __name__ == "__main__": unittest.main()