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.", ) 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": "[]", "competitor_part_numbers": "[]", "series": "[]", "pitch": "[]", "pin_count": "[]", "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 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_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()