import importlib.util import json import os import pathlib import sys import unittest from unittest import mock MODULE_PATH = pathlib.Path(__file__).with_name("artifact_fetch.py") spec = importlib.util.spec_from_file_location("artifact_fetch", MODULE_PATH) assert spec is not None and spec.loader is not None artifact_fetch = importlib.util.module_from_spec(spec) spec.loader.exec_module(artifact_fetch) class DingTalkOnlineDocumentExportTest(unittest.TestCase): def test_text_validation_tolerates_multibyte_char_cut_at_read_boundary(self): # A valid UTF-8 document whose 4096-byte sniff boundary lands in the # middle of a multibyte character must still pass validate_magic. header = "# 标题\n".encode("utf-8") # 7 bytes payload = header + b"a" * (4096 - len(header) - 1) + "机待\u2705".encode("utf-8") + b"\n" self.assertEqual(payload[4095:4098], b"\xe6\x9c\xba") # a char straddles the sniff boundary temp = pathlib.Path("/tmp/kg-md-boundary.md") temp.write_bytes(payload) try: artifact_fetch.validate_magic(str(temp), "md") finally: temp.unlink() def test_text_validation_still_rejects_genuinely_invalid_bytes(self): temp = pathlib.Path("/tmp/kg-md-invalid.md") temp.write_bytes(b"# \xe6\xa0\x87\xe9\xa2\x98\n\xff\xfe not utf-8\n" + b"a" * 8192) try: with self.assertRaises(artifact_fetch.FetchError) as raised: artifact_fetch.validate_magic(str(temp), "md") self.assertEqual(raised.exception.code, "invalid_format") finally: temp.unlink() def test_adoc_uses_export_job_instead_of_ordinary_file_download(self): calls = [] def api(url, headers, payload=None, extra_headers=None, method=None): calls.append((url, payload)) if "/dentry/info" in url: return {"data": {"docKey": "doc-key", "dentryKey": "dentry-key", "extension": "adoc"}} if "/api/document/data" in url: return {"data": {"documentContent": {"checkpoint": {"content": "document-json"}}}} if "/upload_info" in url: return {"data": {"storagePath": "tmp/path", "uploadUrl": "https://oss.example/upload"}} if "/submitExportJob" in url: return {"data": {"jobId": "job-1"}} if "/queryExportJobInfo" in url: return {"data": {"status": "success", "ossUrl": "https://oss.example/result.docx"}} self.fail("unexpected API " + url) with mock.patch.object(artifact_fetch, "dingtalk_json_api", side_effect=api), \ mock.patch.object(artifact_fetch, "dingtalk_put_bytes") as put, \ mock.patch.object(artifact_fetch, "http_get", return_value=("/tmp/result", 123, "a" * 64)) as download: result = artifact_fetch.live_export_dingtalk_document( "raw-node", {"a-token": "x", "x-xsrf-token": "y", "cookie": "z", "referer": "https://alidocs.dingtalk.com/"}, {"name": "说明.adoc", "relative_path": "说明.adoc"}, 1_000_000, ) self.assertEqual(result, ("/tmp/result", 123, "a" * 64, False)) self.assertTrue(any("submitExportJob" in url for url, _ in calls)) self.assertFalse(any("/file/download" in url for url, _ in calls)) put.assert_called_once() download.assert_called_once_with("https://oss.example/result.docx", {}, 180, 1_000_000) class DingTalkDwsFetcherTest(unittest.TestCase): """dws 取件器：路由、回落语义、以及"结果 JSON 形状不得变化"的契约锁定。""" def setUp(self): self.manifest = { "spec_id": "spec-1", "connection_id": "conn-1", "node_id": "sn-test", "source_type": "dingtalk", "source_name": "公司级技术支持知识库", "name": "说明.adoc", "relative_path": "公司级技术支持知识库/说明.adoc", "declared_type": "adoc", "size_bytes": 1000, "modified_at": None, "locator_hash": "loc", } self.ok_verdict = {"ok": True, "code": "", "retryable": False, "missing": False, "error": ""} def _fail(self, code, retryable, missing=False): return {"ok": False, "code": code, "retryable": retryable, "missing": missing, "error": code} def test_dws_adoc_export_is_preferred_and_reported(self): def fake_run(args, timeout=300, cwd=None): self.assertEqual(args[:2], ["doc", "export"]) self.assertIn("--export-format", args) pathlib.Path(cwd, "out.docx").write_bytes(b"PK\x03\x04" + b"x" * 100) return {"success": True}, self.ok_verdict with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "dws"), \ mock.patch.object(artifact_fetch, "dingtalk_raw_keys", return_value=("raw-1", "adoc")), \ mock.patch.object(artifact_fetch, "dws_run", side_effect=fake_run), \ mock.patch.object(artifact_fetch, "exported_doc_is_blank", return_value=False): path, size, digest, matched, fetcher = artifact_fetch.dingtalk_download(self.manifest, 1_000_000) self.assertEqual(fetcher, "dws") self.assertFalse(matched) self.assertEqual(size, 104) self.assertEqual(digest, artifact_fetch.sha256_file(path)) def test_axls_is_routed_through_the_converter(self): built = {} def build(sheets, out): built["sheets"] = sheets pathlib.Path(out).write_bytes(b"PK\x03\x04" + b"y" * 60) return {"sheets": ["S1"]} with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "dws"), \ mock.patch.object(artifact_fetch, "dingtalk_raw_keys", return_value=("raw-2", "axls")), \ mock.patch.object(artifact_fetch.axls_to_xlsx, "fetch_sheets", return_value=[{"name": "S1", "cells": []}]), \ mock.patch.object(artifact_fetch.axls_to_xlsx, "build_xlsx", side_effect=build), \ mock.patch.object(artifact_fetch, "dws_run") as dws_run, \ mock.patch.object(artifact_fetch, "exported_doc_is_blank", return_value=False): path, size, digest, matched, fetcher = artifact_fetch.dingtalk_download( {**self.manifest, "name": "汇总.axls", "declared_type": "axls"}, 1_000_000) self.assertEqual(fetcher, "dws") self.assertTrue(path.endswith("out.xlsx")) self.assertFalse(dws_run.called) # 转换器内部自己调 dws，取件器不再直接调 self.assertEqual(built["sheets"], [{"name": "S1", "cells": []}]) def test_auto_falls_back_to_legacy_on_retryable_failure(self): with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "auto"), \ mock.patch.object(artifact_fetch, "dingtalk_dws_download", side_effect=artifact_fetch.FetchError("dws missing", retryable=True, code="dws_missing")), \ mock.patch.object(artifact_fetch, "dingtalk_local_export", return_value=("/tmp/legacy.docx", True)), \ mock.patch.object(artifact_fetch, "exported_doc_is_blank", return_value=False), \ mock.patch.object(artifact_fetch, "sha256_file", return_value="b" * 64), \ mock.patch.object(artifact_fetch.os.path, "getsize", return_value=42): path, size, digest, matched, fetcher = artifact_fetch.dingtalk_download(self.manifest, 1_000_000) self.assertEqual((path, size, matched, fetcher), ("/tmp/legacy.docx", 42, True, "legacy")) def test_auto_does_not_fall_back_on_terminal_failure(self): with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "auto"), \ mock.patch.object(artifact_fetch, "dingtalk_dws_download", side_effect=artifact_fetch.FetchError("nope", retryable=False, code="forbidden")), \ mock.patch.object(artifact_fetch, "dingtalk_local_export") as legacy: with self.assertRaises(artifact_fetch.FetchError) as raised: artifact_fetch.dingtalk_download(self.manifest, 1_000_000) self.assertEqual(raised.exception.code, "forbidden") self.assertFalse(legacy.called) # 终态错误不得回落（否则掩盖真因） def test_dws_mode_never_falls_back(self): with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "dws"), \ mock.patch.object(artifact_fetch, "dingtalk_dws_download", side_effect=artifact_fetch.FetchError("dws missing", retryable=True, code="dws_missing")), \ mock.patch.object(artifact_fetch, "dingtalk_local_export") as legacy: with self.assertRaises(artifact_fetch.FetchError) as raised: artifact_fetch.dingtalk_download(self.manifest, 1_000_000) self.assertEqual(raised.exception.code, "dws_missing") self.assertFalse(legacy.called) def test_legacy_mode_skips_dws_entirely(self): with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "legacy"), \ mock.patch.object(artifact_fetch, "dingtalk_dws_download") as dws_download, \ mock.patch.object(artifact_fetch, "dingtalk_local_export", return_value=("/tmp/legacy.docx", True)), \ mock.patch.object(artifact_fetch, "exported_doc_is_blank", return_value=False), \ mock.patch.object(artifact_fetch, "sha256_file", return_value="c" * 64), \ mock.patch.object(artifact_fetch.os.path, "getsize", return_value=7): path, size, digest, matched, fetcher = artifact_fetch.dingtalk_download(self.manifest, 1_000_000) self.assertEqual(fetcher, "legacy") self.assertFalse(dws_download.called) def test_catalog_size_over_limit_never_starts_a_download(self): with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "dws"), \ mock.patch.object(artifact_fetch, "dingtalk_raw_keys") as raw_keys, \ mock.patch.object(artifact_fetch, "dws_run") as dws_run: with self.assertRaises(artifact_fetch.FetchError) as raised: artifact_fetch.dingtalk_download({**self.manifest, "size_bytes": 500_000_000}, 1_000_000) self.assertEqual(raised.exception.code, "over_limit") self.assertFalse(raw_keys.called) self.assertFalse(dws_run.called) def test_link_and_mind_nodes_map_to_missing_and_unsupported(self): cases = (("dlink", "链接节点.dlink", "missing", True), ("amind", "培训大纲.amind", "unsupported_format", False)) for declared, name, expected, expected_missing in cases: with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "dws"), \ mock.patch.object(artifact_fetch, "dingtalk_raw_keys", return_value=("raw-3", "file")), \ mock.patch.object(artifact_fetch, "dws_run") as dws_run: with self.assertRaises((artifact_fetch.FetchError, artifact_fetch.MissingSource)) as raised: artifact_fetch.dingtalk_download( {**self.manifest, "name": name, "declared_type": declared}, 1_000_000) if expected_missing: self.assertIsInstance(raised.exception, artifact_fetch.MissingSource) else: self.assertEqual(raised.exception.code, expected) self.assertFalse(dws_run.called) def test_routing_uses_name_suffix_without_declared_type(self): # 生产 manifest 里没有 declared_type 字段；必须靠文件名后缀路由（与老路径一致） def fake_run(args, timeout=300, cwd=None): self.assertEqual(args[:2], ["doc", "export"]) pathlib.Path(cwd, "out.docx").write_bytes(b"PK\x03\x04" + b"q" * 20) return {"success": True}, self.ok_verdict manifest = {k: v for k, v in self.manifest.items() if k != "declared_type"} with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "dws"), \ mock.patch.object(artifact_fetch, "dingtalk_raw_keys", return_value=("raw-9", "file")), \ mock.patch.object(artifact_fetch, "dws_run", side_effect=fake_run), \ mock.patch.object(artifact_fetch, "exported_doc_is_blank", return_value=False): path, size, digest, matched, fetcher = artifact_fetch.dingtalk_download(manifest, 1_000_000) self.assertEqual(fetcher, "dws") self.assertTrue(path.endswith("out.docx")) def test_routing_falls_back_to_dws_inspect_when_name_has_no_suffix(self): calls = [] def fake_run(args, timeout=300, cwd=None): calls.append(args) if args[:2] == ["drive", "+inspect"]: return {"ok": True, "data": {"data": {"file": {"extension": "adoc", "name": "无后缀节点"}}}}, self.ok_verdict pathlib.Path(cwd, "out.docx").write_bytes(b"PK\x03\x04" + b"w" * 30) return {"success": True}, self.ok_verdict with mock.patch.object(artifact_fetch, "DINGTALK_FETCHER", "dws"), \ mock.patch.object(artifact_fetch, "dingtalk_raw_keys", return_value=("raw-10", "file")), \ mock.patch.object(artifact_fetch, "dws_run", side_effect=fake_run), \ mock.patch.object(artifact_fetch, "exported_doc_is_blank", return_value=False): path, size, digest, matched, fetcher = artifact_fetch.dingtalk_download( {k: v for k, v in self.manifest.items() if k != "declared_type"} | {"name": "无后缀节点"}, 1_000_000) self.assertEqual(calls[0][:2], ["drive", "+inspect"]) self.assertEqual(calls[1][:2], ["doc", "export"]) self.assertEqual(fetcher, "dws") def test_result_json_shape_is_unchanged(self): # 平台 Go runner 用 DisallowUnknownFields 解析 stdout：多一个字段就会让所有取件失败。 captured = [] temp = pathlib.Path("/tmp/kg-artifact-fetch-shape.docx") temp.write_bytes(b"PK\x03\x04" + b"z" * 200) try: with mock.patch.object(artifact_fetch, "emit", side_effect=captured.append), \ mock.patch.object(artifact_fetch, "dingtalk_download", return_value=(str(temp), temp.stat().st_size, artifact_fetch.sha256_file(str(temp)), False, "dws")), \ mock.patch.object(sys, "argv", ["artifact_fetch.py", json.dumps(self.manifest)]), \ mock.patch.dict(os.environ, {"INGESTION_SPOOL_ROOT": "/tmp/kg-shape-spool", "ARTIFACT_FETCH_MAX_BYTES": "100000"}): artifact_fetch.main() finally: temp.unlink(missing_ok=True) self.assertEqual(len(captured), 1) self.assertEqual(set(captured[0].keys()), { "ok", "spec_id", "node_id", "connection_id", "sha256", "size_bytes", "spool_ref", "validated_format", "matched_local"}) if __name__ == "__main__": unittest.main()