whale-fall-code
2bb83f3685
App 端免签名接口（设备注册 + detail + CDN 直链），一体化流水线 下载→抽音→传COS→记录→删本地，AI 作品四档判定，四份文档： README（使用指南）、PROGRESS（批次进度）、RE（逆向全过程）、AI（判定规则）。 Co-Authored-By: Claude Code <noreply@anthropic.com>
336 lines
13 KiB
Python
336 lines
13 KiB
Python
"""pipeline.run_one 的端到端测试。不联网、不碰 COS。
|
||
|
||
网络（fetch_detail / download）和 COS（CosUploader）都打桩，但 ffmpeg 抽取、
|
||
ffprobe 时长闸门、记录写入、本地删除全都真跑——这几步才是会悄悄丢数据的地方。
|
||
测试用的 mp4 由 ffmpeg 现造。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
# 移到 scripts/ tests/ 后，直接运行时脚本目录成了 sys.path[0]，
|
||
# 把项目根加回来，import douyin / pipeline 才找得到。
|
||
import sys as _sys
|
||
from pathlib import Path as _Path
|
||
_sys.path.insert(0, str(_Path(__file__).resolve().parent.parent))
|
||
|
||
import argparse
|
||
import json
|
||
import shutil
|
||
import subprocess
|
||
import tempfile
|
||
import threading
|
||
from pathlib import Path
|
||
|
||
import pipeline
|
||
from douyin.client import Aweme
|
||
from douyin.app import DeviceRegisterError
|
||
from douyin.cos import SizeMismatch
|
||
from douyin.download import DownloadResult
|
||
from douyin.meta import RecordLog
|
||
|
||
DUR = 3.0
|
||
|
||
|
||
def make_mp4(path: Path, seconds: float = DUR) -> None:
|
||
subprocess.run(
|
||
["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
||
"-f", "lavfi", "-i", f"testsrc=duration={seconds}:size=320x240:rate=10",
|
||
"-f", "lavfi", "-i", f"sine=frequency=440:duration={seconds}",
|
||
"-c:v", "libx264", "-c:a", "aac", "-shortest", str(path)],
|
||
check=True, capture_output=True)
|
||
|
||
|
||
def fake_aweme(aweme_id="7000000000000000001"):
|
||
raw = {
|
||
"aweme_id": aweme_id,
|
||
"desc": "测试标题 #方言 #四川话",
|
||
"text_extra": [{"type": 1, "hashtag_name": "方言", "hashtag_id": "9",
|
||
"start": 4, "end": 7}],
|
||
"author": {"nickname": "某作者", "uid": "123", "sec_uid": "MS4w"},
|
||
"statistics": {"digg_count": 10, "comment_count": 2, "collect_count": 1,
|
||
"share_count": 3, "play_count": 999},
|
||
"music": {"title": "原声", "author": "某作者"},
|
||
"video": {"duration": 3050, "width": 320, "height": 240},
|
||
}
|
||
return Aweme(aweme_id=aweme_id, desc=raw["desc"], author="某作者",
|
||
create_time=1582635538, video_urls=["http://x/v.mp4"],
|
||
cover_urls=["http://x/c.jpg"], music_url="http://x/m.mp3",
|
||
uri="v0200fg10000test", raw=raw)
|
||
|
||
|
||
class StubPool:
|
||
def current(self):
|
||
return object(), 0
|
||
|
||
def report_dead(self, gen):
|
||
pass
|
||
|
||
|
||
class StubUploader:
|
||
"""记下传了什么，不真传。"""
|
||
|
||
def __init__(self, fail_on=None):
|
||
self.put_calls = []
|
||
self.fail_on = fail_on
|
||
|
||
def key_for(self, name):
|
||
return f"0005/zh/{name}"
|
||
|
||
def put(self, local, key):
|
||
if self.fail_on and self.fail_on in key:
|
||
raise SizeMismatch(f"{key} 大小不符：COS=1 本地=2")
|
||
size = Path(local).stat().st_size
|
||
self.put_calls.append((key, size))
|
||
return f"cos://bucket/{key}", size
|
||
|
||
|
||
def make_args(work, keep_local=False):
|
||
return argparse.Namespace(work_dir=str(work), bitrate="96k", ffprobe="ffprobe",
|
||
keep_local=keep_local)
|
||
|
||
|
||
def run(tmp, *, src_mp4, keep_local=False, uploader=None, aweme=None,
|
||
patch_download=None):
|
||
"""跑一次 run_one，返回 (records, failures, uploader, work_dir)。"""
|
||
work = Path(tmp) / "work"
|
||
work.mkdir(exist_ok=True)
|
||
aweme = aweme or fake_aweme()
|
||
up = uploader or StubUploader()
|
||
log = RecordLog(Path(tmp) / "records.jsonl")
|
||
fails = []
|
||
|
||
def dl(session, urls, dest, **kw):
|
||
shutil.copy(src_mp4, dest)
|
||
return DownloadResult(dest, dest.stat().st_size, resumed=False)
|
||
|
||
orig_dl, orig_fetch = pipeline.download, pipeline.AppClient
|
||
pipeline.download = patch_download or dl
|
||
pipeline.AppClient = lambda dev: type("C", (), {"fetch_detail": lambda s, i: aweme})()
|
||
try:
|
||
pipeline.run_one(
|
||
aweme.aweme_id, "http://src", pool=StubPool(), up=up, log=log,
|
||
args=make_args(work, keep_local), ffmpeg="ffmpeg",
|
||
ff_sem=threading.Semaphore(2), stats=pipeline.Stats(),
|
||
fail=lambda i, s, m: fails.append((i, s, m)),
|
||
abort=threading.Event())
|
||
finally:
|
||
pipeline.download, pipeline.AppClient = orig_dl, orig_fetch
|
||
|
||
recs = [json.loads(l) for l in
|
||
(Path(tmp) / "records.jsonl").read_text(encoding="utf-8").splitlines() if l.strip()] \
|
||
if (Path(tmp) / "records.jsonl").exists() else []
|
||
return recs, fails, up, work
|
||
|
||
|
||
def test_happy_path_writes_record_and_deletes_local():
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = Path(tmp) / "src.mp4"
|
||
make_mp4(src)
|
||
recs, fails, up, work = run(tmp, src_mp4=src)
|
||
|
||
assert not fails, fails
|
||
assert len(recs) == 1, recs
|
||
r = recs[0]
|
||
assert r["aweme_id"] == "7000000000000000001"
|
||
assert r["title"] == "测试标题", repr(r["title"])
|
||
assert r["tags"] == ["方言", "四川话"], r["tags"]
|
||
assert r["author"]["uid"] == "123"
|
||
assert r["stats"]["play"] == 999
|
||
assert abs(r["duration_sec"] - DUR) < 0.5, r["duration_sec"]
|
||
assert r["duration_ms_api"] == 3050 # API 声明值原样留着
|
||
assert r["cos_mp4"].endswith(".mp4") and r["cos_mp3"].endswith(".mp3")
|
||
assert r["raw"]["author"]["sec_uid"] == "MS4w"
|
||
|
||
# 两个都传了
|
||
assert len(up.put_calls) == 2, up.put_calls
|
||
# 本地已删干净
|
||
assert list(work.iterdir()) == [], list(work.iterdir())
|
||
|
||
|
||
def test_keep_local_keeps_files():
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = Path(tmp) / "src.mp4"
|
||
make_mp4(src)
|
||
recs, fails, up, work = run(tmp, src_mp4=src, keep_local=True)
|
||
assert not fails, fails
|
||
names = sorted(p.suffix for p in work.iterdir())
|
||
assert names == [".mp3", ".mp4"], names
|
||
|
||
|
||
def test_upload_failure_keeps_local_and_writes_no_record():
|
||
"""上传失败必须保留本地文件——不然重跑就得重新下载。"""
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = Path(tmp) / "src.mp4"
|
||
make_mp4(src)
|
||
recs, fails, up, work = run(tmp, src_mp4=src,
|
||
uploader=StubUploader(fail_on=".mp3"))
|
||
assert recs == [], recs
|
||
assert len(fails) == 1 and fails[0][1] == "upload", fails
|
||
kept = sorted(p.suffix for p in work.iterdir())
|
||
assert kept == [".mp3", ".mp4"], kept
|
||
|
||
|
||
def test_truncated_mp3_is_caught_before_delete():
|
||
"""时长闸门：mp3 被截断时必须拦下，且不删本地、不写记录。"""
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = Path(tmp) / "src.mp4"
|
||
make_mp4(src, seconds=10.0)
|
||
|
||
orig = pipeline.extract_mp3
|
||
|
||
def truncating(ffmpeg, s, dst, bitrate, timeout=600.0):
|
||
orig(ffmpeg, s, dst, bitrate, timeout)
|
||
data = dst.read_bytes()
|
||
dst.write_bytes(data[: len(data) // 5]) # 砍掉 80%
|
||
|
||
pipeline.extract_mp3 = truncating
|
||
try:
|
||
recs, fails, up, work = run(tmp, src_mp4=src)
|
||
finally:
|
||
pipeline.extract_mp3 = orig
|
||
|
||
assert recs == [], recs
|
||
assert len(fails) == 1 and fails[0][1] == "probe", fails
|
||
assert "截断" in fails[0][2], fails[0][2]
|
||
assert up.put_calls == [], "截断的文件绝不该被上传"
|
||
assert len(list(work.iterdir())) == 2 # 本地保留待查
|
||
|
||
|
||
def test_register_exhausted_aborts_batch():
|
||
"""IP 级注册配额耗尽时必须中止整批，不能把剩余任务逐条烧成失败。
|
||
上一批就是这么烧掉 678 条的。"""
|
||
class DeadPool:
|
||
def current(self):
|
||
raise DeviceRegisterError("所有 host 都返回 device_id=0")
|
||
|
||
def report_dead(self, gen):
|
||
pass
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
work = Path(tmp) / "work"
|
||
work.mkdir()
|
||
log = RecordLog(Path(tmp) / "r.jsonl")
|
||
fails = []
|
||
abort = threading.Event()
|
||
|
||
for i in range(3):
|
||
pipeline.run_one(f"700000000000000000{i}", "u", pool=DeadPool(),
|
||
up=StubUploader(), log=log, args=make_args(work),
|
||
ffmpeg="ffmpeg", ff_sem=threading.Semaphore(2),
|
||
stats=pipeline.Stats(), fail=lambda a, b, c: fails.append((a, b, c)),
|
||
abort=abort)
|
||
|
||
assert abort.is_set(), "第一条就该拉起中止信号"
|
||
# 只有触发中止的那条记一次失败，后两条直接返回，不重复撞注册
|
||
assert len(fails) == 1, fails
|
||
assert fails[0][1] == "abort", fails
|
||
|
||
|
||
def test_download_failure_is_reported():
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = Path(tmp) / "src.mp4"
|
||
make_mp4(src)
|
||
|
||
def boom(session, urls, dest, **kw):
|
||
raise OSError("连接被重置")
|
||
|
||
recs, fails, up, work = run(tmp, src_mp4=src, patch_download=boom)
|
||
assert recs == [] and len(fails) == 1
|
||
assert fails[0][1] == "download" and "连接被重置" in fails[0][2], fails
|
||
|
||
|
||
def test_resume_skips_completed():
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = Path(tmp) / "src.mp4"
|
||
make_mp4(src)
|
||
run(tmp, src_mp4=src)
|
||
recs2, fails2, up2, _ = run(tmp, src_mp4=src) # 再跑一次同一条
|
||
assert len(recs2) == 1, "重复跑不该写出第二行"
|
||
assert not fails2, fails2
|
||
|
||
|
||
def test_sweep_removes_files_already_in_records():
|
||
"""work/ 里凡是已经在 records.jsonl 里的文件都是废物：早就传上去了，
|
||
留着只会撑爆磁盘。实测一次 66411 条的批跑到 2.7% 就攒了 11 GB。
|
||
但**失败项的文件必须留下**——那是重跑续传的本钱，见
|
||
test_upload_failure_keeps_local_and_writes_no_record。"""
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
work = Path(tmp) / "work"
|
||
work.mkdir()
|
||
recs = Path(tmp) / "records.jsonl"
|
||
recs.write_text(
|
||
json.dumps({"aweme_id": "111", "cos_mp4": "cos://x/111.mp4"}) + chr(10)
|
||
+ json.dumps({"aweme_id": "222", "cos_mp4": "cos://x/222.mp4"}) + chr(10),
|
||
encoding="utf-8")
|
||
for name in ("111.mp4", "111.mp3", "222.mp4", "333.mp4", "333.mp3"):
|
||
(work / name).write_bytes(b"x" * 10)
|
||
(work / "444.mp4.part").write_bytes(b"x" * 10)
|
||
|
||
freed = pipeline.sweep_work_dir(work, recs)
|
||
|
||
left = sorted(p.name for p in work.iterdir())
|
||
assert left == ["333.mp3", "333.mp4"], left # 333 没入库，留着重跑
|
||
# 111.mp4 + 111.mp3 + 222.mp4 已入库，加上 444.mp4.part
|
||
assert freed == 4, freed
|
||
|
||
|
||
def test_sweep_tolerates_missing_paths():
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
work = Path(tmp) / "work"
|
||
assert pipeline.sweep_work_dir(work, Path(tmp) / "nope.jsonl") == 0
|
||
work.mkdir()
|
||
assert pipeline.sweep_work_dir(work, Path(tmp) / "nope.jsonl") == 0
|
||
|
||
|
||
def test_plan_batch_splits_long_runs():
|
||
"""长跑必须切段：单个进程跑到 1.5 万条上下，Windows 会话的桌面堆就被
|
||
几万个 ffmpeg/ffprobe 子进程耗光，之后 CreateProcess 一律失败
|
||
（rc=3221225794 / STATUS_DLL_INIT_FAILED）。表现是完成数冻结、失败数
|
||
猛涨、work/ 暴涨。实测两次：三进程并跑时撞过一次，单进程跑到 16399 条
|
||
又撞一次 —— 所以是累积量的问题，不是并发数。切段后每段跑完重启进程，
|
||
桌面堆跟着重置。"""
|
||
pending = [(str(i), f"u{i}") for i in range(100)]
|
||
batch, more = pipeline.plan_batch(pending, 30)
|
||
assert len(batch) == 30 and more is True
|
||
batch, more = pipeline.plan_batch(pending, 100)
|
||
assert len(batch) == 100 and more is False, "刚好跑完不该再重启"
|
||
batch, more = pipeline.plan_batch(pending, 500)
|
||
assert len(batch) == 100 and more is False
|
||
batch, more = pipeline.plan_batch(pending, 0)
|
||
assert len(batch) == 100 and more is False, "0 = 不切段"
|
||
|
||
|
||
def test_next_abort_wait_backs_off():
|
||
"""撞 IP 级注册限流后要等冷却再续跑。连着撞说明还没解封，等得越来越久；
|
||
一旦有一段跑成了就把等待打回底价。上限别太长，不然白白空转。"""
|
||
assert pipeline.next_abort_wait(0, base=1800, cap=14400) == 1800
|
||
assert pipeline.next_abort_wait(1, base=1800, cap=14400) == 3600
|
||
assert pipeline.next_abort_wait(2, base=1800, cap=14400) == 7200
|
||
assert pipeline.next_abort_wait(3, base=1800, cap=14400) == 14400
|
||
assert pipeline.next_abort_wait(9, base=1800, cap=14400) == 14400, "要封顶"
|
||
assert pipeline.next_abort_wait(0, base=0, cap=14400) == 0, "0=不等，直接退出"
|
||
|
||
|
||
def main() -> int:
|
||
if not shutil.which("ffmpeg") or not shutil.which("ffprobe"):
|
||
print("跳过：需要 ffmpeg/ffprobe")
|
||
return 0
|
||
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
||
failed = 0
|
||
for t in tests:
|
||
try:
|
||
t()
|
||
print(f" ok {t.__name__}")
|
||
except AssertionError as exc:
|
||
failed += 1
|
||
print(f" FAIL {t.__name__}: {exc}")
|
||
except Exception as exc:
|
||
failed += 1
|
||
print(f" ERR {t.__name__}: {type(exc).__name__}: {exc}")
|
||
print(f"\n{len(tests) - failed}/{len(tests)} 通过")
|
||
return 1 if failed else 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|