import hashlib import json import tempfile import unittest from pathlib import Path from unittest.mock import patch from obsidian_wiki.access_control import source_decision as actual_source_decision 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, mcp_runtime_options, read_file_asset_resource, record_allowed_for_mcp, 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._write_profile_file() self.drawing_id = self._write_asset_file("drawings/a.pdf", b"%PDF allowed drawing") self.order_id = self._write_asset_file("orders/order.xlsx", b"customer price data") self.complaint_id = 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"], allowed_bots=["bjd-general-bot", "missing-bot"], ) 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) return "sha256:" + hashlib.sha256(data).hexdigest() 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, }, "public-product-bot": { "allowed_visibility": ["public", "internal"], "allowed_data_classification": [ "public", "internal", "technical-confidential", ], "allowed_business_domains": ["products", "engineering"], "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): shareable_id = 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/" + self.drawing_id.removeprefix("sha256:"), ) 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, allow_source_files=True, ) 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, allow_source_files=True, ) 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, allow_source_files=True, ) def test_resource_read_requires_explicit_server_opt_in(self): with self.assertRaisesRegex(AssetAccessError, "source file reads are disabled"): read_file_asset_resource( self.vault, asset_uri_for_file_id(self.drawing_id), db_path=self.db_path, ) def test_runtime_options_source_file_gate_defaults_closed_and_reads_env(self): with patch.dict("os.environ", {}, clear=True): self.assertFalse(mcp_runtime_options()["allow_source_files"]) with patch.dict("os.environ", {"OBSIDIAN_MCP_ALLOW_SOURCE_FILES": "1"}, clear=True): self.assertTrue(mcp_runtime_options()["allow_source_files"]) def test_resource_read_rejects_db_registry_archive_path_mismatch(self): other_path = "_sources/archive/drawings/other.pdf" self._write_asset_file("drawings/other.pdf", b"%PDF other drawing") with connect(self.db_path) as conn: conn.execute( "UPDATE file_assets SET archive_path = ? WHERE file_id = ?", (other_path, self.drawing_id), ) with self.assertRaisesRegex(AssetAccessError, "archive_path mismatch"): read_file_asset_resource( self.vault, asset_uri_for_file_id(self.drawing_id), db_path=self.db_path, allow_source_files=True, ) def test_resource_read_accepts_normalization_equivalent_archive_paths(self): with connect(self.db_path) as conn: conn.execute( "UPDATE file_assets SET archive_path = ? WHERE file_id = ?", ("_sources/archive/drawings/../drawings/a.pdf", self.drawing_id), ) data = read_file_asset_resource( self.vault, asset_uri_for_file_id(self.drawing_id), db_path=self.db_path, allow_source_files=True, ) self.assertEqual(data["bytes"], b"%PDF allowed drawing") def test_resource_read_rejects_content_that_no_longer_matches_file_id(self): (self.vault / "_sources" / "archive" / "drawings" / "a.pdf").write_bytes( b"%PDF tampered drawing" ) with self.assertRaisesRegex(AssetAccessError, "content hash mismatch"): read_file_asset_resource( self.vault, asset_uri_for_file_id(self.drawing_id), db_path=self.db_path, allow_source_files=True, ) 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_with_valid_profile(self): with self.assertRaises(AssetAccessError): get_file_asset_metadata( self.vault, self.drawing_id, db_path=self.db_path, profile_name="public-product-bot", ) def test_missing_requested_profile_fails_closed(self): with self.assertRaisesRegex(AssetAccessError, "bot profile not found"): search_file_assets( self.vault, "B3501", db_path=self.db_path, profile_name="missing-bot", ) with self.assertRaisesRegex(AssetAccessError, "bot profile not found"): get_file_asset_metadata( self.vault, self.drawing_id, db_path=self.db_path, profile_name="missing-bot", ) def test_metadata_acl_does_not_rewrite_profile_source_file_setting(self): profile = { "allowed_visibility": ["internal"], "allowed_data_classification": ["technical-confidential"], "allowed_business_domains": ["products", "engineering"], "include_source_files": False, } record = { "archive_path": "_sources/archive/drawings/a.pdf", "doc_type": "drawing", "visibility": "internal", "data_classification": "customer-shareable", "business_categories": ["products/btb", "drawings/customer"], "allowed_bots": ["bjd-general-bot"], } with patch( "obsidian_wiki.mcp_assets.source_decision", wraps=actual_source_decision, ) as mocked_decision: allowed, reason = record_allowed_for_mcp( record, profile_name="bjd-general-bot", profile=profile, ) self.assertTrue(allowed, reason) runtime_profile = mocked_decision.call_args.args[1] self.assertIsNot(runtime_profile, profile) self.assertTrue(runtime_profile["include_source_files"]) self.assertIn("customer-shareable", runtime_profile["allowed_data_classification"]) self.assertFalse(profile["include_source_files"]) self.assertNotIn("customer-shareable", profile["allowed_data_classification"]) def test_search_filters_acl_before_applying_limit(self): allowed_id = self._write_asset_file( "drawings/allowed-limit.pdf", b"%PDF allowed lower-ranked limit result", ) self._insert_asset( allowed_id, archive_path="_sources/archive/drawings/allowed-limit.pdf", original_name="Allowed limit drawing.pdf", doc_type="drawing", source_type="pdf", product_lines=["BTB"], business_categories=["products/btb"], data_classification="technical-confidential", visibility="internal", summary="BLOCKME appears only as unstructured text.", ) for index in range(4): blocked_id = self._write_asset_file( f"drawings/blocked-limit-{index}.pdf", f"%PDF blocked limit result {index}".encode(), ) self._insert_asset( blocked_id, archive_path=f"_sources/archive/drawings/blocked-limit-{index}.pdf", original_name=f"Blocked limit drawing {index}.pdf", doc_type="drawing", source_type="pdf", product_lines=["BTB"], business_categories=["products/btb"], data_classification="technical-confidential", visibility="internal", summary="BLOCKME structured identifier result.", bjd_part_numbers=["BLOCKME"], allowed_bots=["other-bot"], ) rows = search_file_assets(self.vault, "BLOCKME", db_path=self.db_path, limit=1) self.assertEqual([row["file_id"] for row in rows], [allowed_id]) def test_search_clamps_limit_to_public_range(self): for index in range(101): file_id = self._write_asset_file( f"drawings/bulk-{index}.pdf", f"%PDF BULK-LIMIT {index}".encode(), ) self._insert_asset( file_id, archive_path=f"_sources/archive/drawings/bulk-{index}.pdf", original_name=f"Bulk limit drawing {index}.pdf", doc_type="drawing", source_type="pdf", product_lines=["BTB"], business_categories=["products/btb"], data_classification="technical-confidential", visibility="internal", summary="BULK-LIMIT clamp candidate.", ) lower = search_file_assets(self.vault, "B3501", db_path=self.db_path, limit=0) upper = search_file_assets(self.vault, "BULK-LIMIT", db_path=self.db_path, limit=1000) self.assertEqual(len(lower), 1) self.assertEqual(len(upper), 100) if __name__ == "__main__": unittest.main()