167 lines
8.3 KiB
Python
167 lines
8.3 KiB
Python
"""Real FFmpeg + native messaging checks. Add --online for a small public YouTube download."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import io
|
|
import json
|
|
import queue
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
from host.common import EXTENSION_ID
|
|
from host.media import NO_WINDOW, capabilities, config, probe, tool
|
|
from host.protocol import read_message, write_message
|
|
from host.worker import convert
|
|
|
|
|
|
class NativeSession:
|
|
def __init__(self):
|
|
launcher = ROOT / ".runtime" / "launch-host.cmd"
|
|
self.process = subprocess.Popen([str(launcher), f"chrome-extension://{EXTENSION_ID}/"],
|
|
stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
creationflags=NO_WINDOW)
|
|
self.messages = queue.Queue()
|
|
self.jobs = {}
|
|
self.errors = []
|
|
def read():
|
|
try:
|
|
while (message := read_message(self.process.stdout)) is not None:
|
|
self.messages.put(message)
|
|
except Exception as error:
|
|
self.messages.put({"protocolError": str(error)})
|
|
threading.Thread(target=read, daemon=True).start()
|
|
threading.Thread(target=lambda: self.errors.append(self.process.stderr.read().decode("utf-8", "replace")), daemon=True).start()
|
|
|
|
def send(self, action, payload=None):
|
|
token = str(uuid.uuid4())
|
|
write_message(self.process.stdin, {"id": token, "action": action, "payload": payload or {}})
|
|
return token
|
|
|
|
def next(self, timeout=150):
|
|
message = self.messages.get(timeout=timeout)
|
|
if "protocolError" in message:
|
|
raise AssertionError(message)
|
|
if message.get("event") == "job":
|
|
self.jobs[message["job"]["id"]] = message["job"]
|
|
return message
|
|
|
|
def call(self, action, payload=None):
|
|
token = self.send(action, payload)
|
|
while True:
|
|
message = self.next()
|
|
if message.get("id") == token:
|
|
if not message["ok"]:
|
|
raise AssertionError(message["error"])
|
|
return message["result"]
|
|
|
|
def download(self, payload):
|
|
result = self.call("download", payload)
|
|
job_id = result["jobId"]
|
|
stage = None
|
|
deadline = time.monotonic() + 240
|
|
while time.monotonic() < deadline:
|
|
job = self.jobs.get(job_id, {})
|
|
if job.get("status") == "complete": return job["result"]
|
|
if job.get("status") in {"error", "cancelled"}: raise AssertionError(job)
|
|
self.next()
|
|
job = self.jobs.get(job_id, {})
|
|
if job.get("stage") != stage:
|
|
stage = job.get("stage"); print(f"Native job: {stage}", flush=True)
|
|
raise TimeoutError("Native job timeout")
|
|
|
|
def close(self):
|
|
self.process.stdin.close()
|
|
self.process.wait(timeout=30)
|
|
if self.process.returncode:
|
|
raise AssertionError("Native host exit: " + "".join(self.errors))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--online", action="store_true")
|
|
args = parser.parse_args()
|
|
artifacts = ROOT / "test-artifacts"
|
|
artifacts.mkdir(exist_ok=True)
|
|
media_dir = artifacts / "media"
|
|
media_dir.mkdir(exist_ok=True)
|
|
hardware = capabilities()
|
|
synthetic = media_dir / "input-vp9-opus.webm"
|
|
subprocess.run([tool("ffmpeg"), "-hide_banner", "-loglevel", "error", "-y", "-f", "lavfi", "-i", "testsrc2=s=640x360:r=30",
|
|
"-f", "lavfi", "-i", "sine=frequency=440:sample_rate=48000", "-t", "0.6", "-c:v", "libvpx-vp9",
|
|
"-cpu-used", "8", "-c:a", "libopus", str(synthetic)], check=True, creationflags=NO_WINDOW)
|
|
summary = {"gpu": hardware["gpu"], "nvenc": hardware["nvenc"], "formats": [], "native": {}}
|
|
defaults = {"mode": "video", "container": "mp4", "codec": "h264", "engine": "nvenc", "preset": "balanced", "bitrate": 256}
|
|
source_probe = probe(synthetic)
|
|
for codec in ("h264", "hevc", "av1"):
|
|
target = media_dir / f"gpu-{codec}.mp4"
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
plan = convert(synthetic, target, {**defaults, "codec": codec}, source_probe, hardware["nvenc"])
|
|
streams = probe(target)["streams"]
|
|
assert plan["engine"] == "nvenc"
|
|
assert {s["codec_type"] for s in streams} >= {"video", "audio"}
|
|
assert next(s for s in streams if s["codec_type"] == "video")["codec_name"] == codec
|
|
summary["formats"].append({"format": codec, "engine": plan["engine"], "size": target.stat().st_size})
|
|
for container in ("m4a", "mp3", "opus", "flac", "wav"):
|
|
target = media_dir / f"audio.{container}"
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
convert(synthetic, target, {**defaults, "mode": "audio", "container": container}, source_probe, hardware["nvenc"])
|
|
assert {s["codec_type"] for s in probe(target)["streams"]} == {"audio"}
|
|
summary["formats"].append({"format": container, "size": target.stat().st_size})
|
|
target = media_dir / "video-only.mkv"
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
plan = convert(synthetic, target, {**defaults, "mode": "video_only", "container": "mkv", "codec": "auto"}, source_probe, hardware["nvenc"])
|
|
assert plan["engine"] == "copy"
|
|
assert {s["codec_type"] for s in probe(target)["streams"]} == {"video"}
|
|
with_session = NativeSession()
|
|
try:
|
|
hello = with_session.call("hello")
|
|
summary["native"]["hello"] = hello["gpu"]
|
|
if args.online:
|
|
url = "https://www.youtube.com/watch?v=jNQXAC9IVRw"
|
|
info = with_session.call("inspect", {"url": url})
|
|
(artifacts / "sample-metadata.json").write_text(json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
separated = [f for f in info["videos"] if not f["hasAudio"]]
|
|
assert separated and info["audios"], "Separate tracks required for this check"
|
|
chosen = sorted(separated, key=lambda f: (f["height"], f["fps"]))[0]
|
|
payload = {**defaults, "url": url, "folder": str(media_dir), "videoId": chosen["id"], "audioId": "best", "filename": "YouTube-test-HEVC", "codec": "hevc"}
|
|
result = with_session.download(payload)
|
|
downloaded = probe(result["path"])
|
|
assert {s["codec_type"] for s in downloaded["streams"]} >= {"video", "audio"}
|
|
assert result["engine"] == "nvenc"
|
|
summary["native"]["videoPlusAudio"] = result
|
|
result = with_session.download({**payload, "mode": "audio", "container": "mp3", "filename": "YouTube-test-audio"})
|
|
assert {s["codec_type"] for s in probe(result["path"])["streams"]} == {"audio"}
|
|
summary["native"]["audioOnly"] = result
|
|
result = with_session.download({**payload, "mode": "video_only", "container": "mkv", "codec": "copy", "filename": "YouTube-test-video-only"})
|
|
assert {s["codec_type"] for s in probe(result["path"])["streams"]} == {"video"}
|
|
summary["native"]["videoOnly"] = result
|
|
long_url = "https://www.youtube.com/watch?v=aircAruvnKk"
|
|
cancel_payload = {**payload, "url": long_url, "videoId": "best", "filename": "cancelled-test"}
|
|
job_id = with_session.call("download", cancel_payload)["jobId"]
|
|
while True:
|
|
message = with_session.next()
|
|
job = with_session.jobs.get(job_id, {})
|
|
if job.get("stage") == "downloading": break
|
|
if job.get("status") == "error": raise AssertionError(job)
|
|
with_session.call("cancel", {"jobId": job_id})
|
|
deadline = time.monotonic() + 8
|
|
while time.monotonic() < deadline and list(media_dir.glob(".ytdl-*")):
|
|
time.sleep(0.1)
|
|
assert not list(media_dir.glob(".ytdl-*")), "Cancelled job left temporary files"
|
|
assert not list(media_dir.glob("cancelled-test.*")), "Cancelled file was published"
|
|
summary["native"]["cancel"] = "passed; process tree and temporary files removed"
|
|
finally:
|
|
with_session.close()
|
|
(artifacts / "verification.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print("Integration checks passed.", flush=True)
|
|
|
|
|
|
if __name__ == "__main__": main()
|