345 lines
15 KiB
Python
345 lines
15 KiB
Python
# -*- 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"] == "内容"
|