姚章浩
257be88ecf
validate_magic read exactly 4096 bytes and strict-decoded them; a multibyte character straddling the boundary made perfectly valid .md/.txt files fail with invalid_format. Only tolerate 'unexpected end of data' when the file actually continues past the sniff window; genuinely corrupt bytes still fail.
71 lines
3.4 KiB
Python
71 lines
3.4 KiB
Python
import importlib.util
|
|
import pathlib
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|