# -*- coding: utf-8 -*- """youtube_export_download.py 测试。 覆盖不依赖真实浏览器/登录态的全部逻辑: dedup_path() 重名 ` (n)` 去重 build_export_filename() 从 exportQuery 反推 zip 文件名 decode_zipped_data() zippedData base64 解码 intercept_and_save() 响应拦截落盘(用 FakePage/FakeResponse 模拟 Playwright) _default_user_data_dir() 浏览器用户数据目录推导 selftest() / --selftest 内置自测 run() 需要已登录 YouTube Studio 的真实浏览器会话,不在自动化测试范围(见 SKILL.md)。 """ import base64 import io import json import zipfile from pathlib import Path import pytest from conftest import DOWNLOADER_SCRIPT, run_script # --------------------------------------------------------------------------- # 测试替身:模拟 playwright Page / Response / Request # --------------------------------------------------------------------------- class FakeRequest: def __init__(self, post_data=None): self.post_data = post_data class FakeResponse: def __init__(self, url, payload, post_data=None): self.url = url self._payload = payload self.request = FakeRequest(post_data) def json(self): return self._payload class FakePage: """记录 page.on() 注册的回调,测试中手动触发。""" def __init__(self): self.handlers = {} def on(self, event, handler): self.handlers[event] = handler # --------------------------------------------------------------------------- # 工具 # --------------------------------------------------------------------------- def make_zip_bytes(files): """构造内存 zip:{成员名: 内容(bytes/str)}。""" buf = io.BytesIO() with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: for name, data in files.items(): zf.writestr(name, data) return buf.getvalue() def make_export_query(dimension="USER", start=20260723, end=20260820, with_time=True): query = {"dimensions": [{"type": dimension}]} if with_time: query["timeRange"] = {"dateIdRange": { "inclusiveStart": start, "exclusiveEnd": end}} return {"joinRequest": {"nodes": [{"value": {"query": query}}]}} def make_export_response(tmp_dir, files=None, account="WL Media", dimension="USER", post_data=None): """构造一个命中 csv_export 的 FakeResponse(zip 为真实可解压内容)。""" data = make_zip_bytes(files or {"表格数据.csv": "a,b\n1,2"}) payload = {"zippedData": base64.b64encode(data).decode("ascii")} if post_data is None: post_data = json.dumps({"exportQuery": make_export_query(dimension=dimension)}) url = "https://studio.youtube.com/youtubei/v1/yta_web/csv_export?alt=json" return FakeResponse(url, payload, post_data), data # --------------------------------------------------------------------------- # dedup_path:重名去重 # --------------------------------------------------------------------------- class TestDedupPath: def test_no_conflict_returns_original(self, downloader, tmp_path): assert Path(downloader.dedup_path(str(tmp_path), "需求文件.zip")).name == "需求文件.zip" def test_conflict_appends_suffix_in_order(self, downloader, tmp_path): d = str(tmp_path) Path(downloader.dedup_path(d, "需求文件.zip")).write_bytes(b"a") # 需求文件.zip p1 = Path(downloader.dedup_path(d, "需求文件.zip")) assert p1.name == "需求文件 (1).zip" p1.write_bytes(b"b") p2 = Path(downloader.dedup_path(d, "需求文件.zip")) assert p2.name == "需求文件 (2).zip" p2.write_bytes(b"c") assert Path(downloader.dedup_path(d, "需求文件.zip")).name == "需求文件 (3).zip" def test_suffix_before_extension(self, downloader, tmp_path): """后缀插在扩展名之前:`名称 (1).zip` 而非 `名称.zip (1)`。""" d = str(tmp_path) Path(downloader.dedup_path(d, "report.2026.zip")).write_bytes(b"x") assert Path(downloader.dedup_path(d, "report.2026.zip")).name == "report.2026 (1).zip" def test_other_names_unaffected(self, downloader, tmp_path): d = str(tmp_path) Path(downloader.dedup_path(d, "A.zip")).write_bytes(b"x") assert Path(downloader.dedup_path(d, "B.zip")).name == "B.zip" def test_gap_filling(self, downloader, tmp_path): """(1) 被删后再次落盘优先复用空位。""" d = str(tmp_path) p0 = Path(downloader.dedup_path(d, "f.zip")); p0.write_bytes(b"0") p1 = Path(downloader.dedup_path(d, "f.zip")); p1.write_bytes(b"1") p1.unlink() # 删掉 (1) assert Path(downloader.dedup_path(d, "f.zip")).name == "f (1).zip" # --------------------------------------------------------------------------- # build_export_filename:文件名反推 # --------------------------------------------------------------------------- class TestBuildExportFilename: def test_user_dimension(self, downloader): name = downloader.build_export_filename(make_export_query(), "WL Media") assert name == "频道 2026-07-23_2026-08-20 WL Media.zip" def test_video_dimension(self, downloader): q = make_export_query(dimension="VIDEO") assert downloader.build_export_filename(q, "acct") == "内容 2026-07-23_2026-08-20 acct.zip" def test_content_owner_dimension(self, downloader): q = make_export_query(dimension="CONTENT_OWNER") assert downloader.build_export_filename(q, "acct") == "内容 2026-07-23_2026-08-20 acct.zip" def test_custom_dimension_label_overrides(self, downloader): q = make_export_query(dimension="VIDEO") name = downloader.build_export_filename(q, "acct", dimension_label="地区") assert name == "地区 2026-07-23_2026-08-20 acct.zip" def test_unknown_dimension_empty_label(self, downloader): q = make_export_query(dimension="SOMETHING") assert downloader.build_export_filename(q, "acct") == " 2026-07-23_2026-08-20 acct.zip" def test_date_id_formatting(self, downloader): """dateId 20260101 -> 2026-01-01(8 位定长切分)。""" q = make_export_query(start=20260101, end=20260102) assert "2026-01-01_2026-01-02" in downloader.build_export_filename(q, "acct") def test_missing_date_range_raises(self, downloader): q = make_export_query(with_time=False) with pytest.raises(ValueError, match="日期范围"): downloader.build_export_filename(q, "acct") def test_empty_query_raises(self, downloader): with pytest.raises(ValueError, match="日期范围"): downloader.build_export_filename({}, "acct") def test_multi_node_query(self, downloader): """维度与日期分布在不同 node 时也能各取所需。""" q = {"joinRequest": {"nodes": [ {"value": {"query": {"dimensions": [{"type": "VIDEO"}]}}}, {"value": {"query": {"timeRange": {"dateIdRange": { "inclusiveStart": 20260723, "exclusiveEnd": 20260820}}}}}, ]}} assert downloader.build_export_filename(q, "acct") == "内容 2026-07-23_2026-08-20 acct.zip" # --------------------------------------------------------------------------- # decode_zipped_data:zippedData 解码 # --------------------------------------------------------------------------- class TestDecodeZippedData: def test_roundtrip(self, downloader): data = make_zip_bytes({"表格数据.csv": "a,b\n1,2"}) payload = {"zippedData": base64.b64encode(data).decode("ascii")} assert downloader.decode_zipped_data(payload) == data def test_decoded_bytes_are_valid_zip(self, downloader): data = make_zip_bytes({"x.csv": "1,2"}) payload = {"zippedData": base64.b64encode(data).decode("ascii")} out = downloader.decode_zipped_data(payload) with zipfile.ZipFile(io.BytesIO(out)) as zf: assert zf.read("x.csv") == b"1,2" def test_missing_field_raises(self, downloader): with pytest.raises(ValueError, match="zippedData"): downloader.decode_zipped_data({"foo": "bar"}) @pytest.mark.parametrize("empty", ["", None]) def test_empty_field_raises(self, downloader, empty): with pytest.raises(ValueError, match="zippedData"): downloader.decode_zipped_data({"zippedData": empty}) # --------------------------------------------------------------------------- # intercept_and_save:拦截器(FakePage 模拟) # --------------------------------------------------------------------------- class TestInterceptAndSave: def test_saves_zip_and_reports(self, downloader, tmp_path): page = FakePage() saved = downloader.intercept_and_save(page, "WL Media", str(tmp_path)) resp, data = make_export_response(tmp_path, {"表格数据.csv": "a,b\n1,2"}) page.handlers["response"](resp) assert len(saved) == 1 filename, size, path = saved[0] assert filename == "频道 2026-07-23_2026-08-20 WL Media.zip" assert size == len(data) assert Path(path).read_bytes() == data # 落盘内容与响应一致 with zipfile.ZipFile(path) as zf: # 且是可解压的有效 zip assert zf.read("表格数据.csv") == b"a,b\n1,2" def test_second_export_deduped(self, downloader, tmp_path): page = FakePage() saved = downloader.intercept_and_save(page, "WL Media", str(tmp_path)) resp, _ = make_export_response(tmp_path) page.handlers["response"](resp) page.handlers["response"](resp) # saved 记录的是请求体反推出的原始文件名;实际落盘路径(s[2])带去重后缀 assert [s[0] for s in saved] == ["频道 2026-07-23_2026-08-20 WL Media.zip"] * 2 assert [Path(s[2]).name for s in saved] == [ "频道 2026-07-23_2026-08-20 WL Media.zip", "频道 2026-07-23_2026-08-20 WL Media (1).zip", ] def test_ignores_non_export_responses(self, downloader, tmp_path): page = FakePage() saved = downloader.intercept_and_save(page, "WL Media", str(tmp_path)) other = FakeResponse("https://studio.youtube.com/youtubei/v1/other", {}) page.handlers["response"](other) assert saved == [] assert list(tmp_path.iterdir()) == [] # 无任何落盘 def test_illegal_chars_sanitized(self, downloader, tmp_path): """Windows 非法字符 \\ / : * ? " < > | 替换为 _。""" page = FakePage() saved = downloader.intercept_and_save(page, r'A/B:C*D?E"FH|I', str(tmp_path)) resp, _ = make_export_response(tmp_path) page.handlers["response"](resp) filename = saved[0][0] assert filename == "频道 2026-07-23_2026-08-20 A_B_C_D_E_F_G_H_I.zip" assert Path(saved[0][2]).exists() def test_bad_post_data_falls_back_to_export_zip(self, downloader, tmp_path): """exportQuery 反推失败时兜底 export.zip,内容不丢。""" page = FakePage() saved = downloader.intercept_and_save(page, "WL Media", str(tmp_path)) resp, data = make_export_response(tmp_path, post_data="not-json-{{{") page.handlers["response"](resp) assert saved[0][0] == "export.zip" assert Path(saved[0][2]).read_bytes() == data def test_empty_post_data_falls_back(self, downloader, tmp_path): """request.post_data 为 None(或 "{}"):反推失败兜底 export.zip。""" page = FakePage() saved = downloader.intercept_and_save(page, "acct", str(tmp_path)) data = make_zip_bytes({"x.csv": "1"}) resp = FakeResponse( "https://studio.youtube.com/youtubei/v1/yta_web/csv_export?alt=json", {"zippedData": base64.b64encode(data).decode("ascii")}, post_data=None, ) page.handlers["response"](resp) assert saved[0][0] == "export.zip" assert Path(saved[0][2]).read_bytes() == data def test_missing_zipped_data_no_save_with_stderr(self, downloader, tmp_path, capsys): page = FakePage() saved = downloader.intercept_and_save(page, "WL Media", str(tmp_path)) resp = FakeResponse( "https://studio.youtube.com/youtubei/v1/yta_web/csv_export?alt=json", {"foo": "bar"}, post_data=json.dumps({"exportQuery": make_export_query()}), ) page.handlers["response"](resp) assert saved == [] assert list(tmp_path.iterdir()) == [] assert "处理 csv_export 响应失败" in capsys.readouterr().err # --------------------------------------------------------------------------- # _default_user_data_dir:用户数据目录推导 # --------------------------------------------------------------------------- class TestDefaultUserDataDir: def test_chrome(self, downloader, monkeypatch): monkeypatch.setenv("LOCALAPPDATA", r"C:\fake\Local") assert downloader._default_user_data_dir("chrome") == \ r"C:\fake\Local\Google\Chrome\User Data" def test_msedge(self, downloader, monkeypatch): monkeypatch.setenv("LOCALAPPDATA", r"C:\fake\Local") assert downloader._default_user_data_dir("msedge") == \ r"C:\fake\Local\Microsoft\Edge\User Data" def test_default_is_chrome(self, downloader, monkeypatch): monkeypatch.setenv("LOCALAPPDATA", r"C:\fake\Local") assert downloader._default_user_data_dir(None) == \ r"C:\fake\Local\Google\Chrome\User Data" def test_fallback_without_env(self, downloader, monkeypatch): monkeypatch.delenv("LOCALAPPDATA", raising=False) expected = str(Path.home() / "AppData" / "Local" / "Google" / "Chrome" / "User Data") assert downloader._default_user_data_dir("chrome") == expected # --------------------------------------------------------------------------- # 内置自测与 CLI # --------------------------------------------------------------------------- class TestSelftest: def test_module_selftest(self, downloader, capsys): downloader.selftest() assert "selftest OK" in capsys.readouterr().out def test_cli_selftest(self): r = run_script(DOWNLOADER_SCRIPT, ["--selftest"]) assert r.returncode == 0 assert "selftest OK" in r.stdout def test_cli_no_args_runs_selftest_and_prints_usage(self): """无参数运行:先自测,再提示实际运行方式。""" r = run_script(DOWNLOADER_SCRIPT, []) assert r.returncode == 0 assert "selftest OK" in r.stdout assert "--user-data-dir" in r.stdout assert "--connect" in r.stdout # --------------------------------------------------------------------------- # 模块常量契约 # --------------------------------------------------------------------------- class TestConstants: def test_csv_export_path(self, downloader): assert downloader.CSV_EXPORT_PATH == "/youtubei/v1/yta_web/csv_export" def test_dimension_labels(self, downloader): assert downloader.DIMENSION_LABEL["VIDEO"] == "内容" assert downloader.DIMENSION_LABEL["USER"] == "频道" assert downloader.DIMENSION_LABEL["CONTENT_OWNER"] == "内容"