feat(scripts): 添加批量生成 YouTube Studio 内容管理器 URL 的脚本

2026-08-21 18:21:03 +08:00
commit c1b39e9466
29 changed files with 4020 additions and 0 deletions
--- a/tests/test_youtube_export_download.py
+++ b/tests/test_youtube_export_download.py
@@ -0,0 +1,344 @@
+# -*- 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"F<G>H|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"] == "内容"