feat: compress downloads to max 250MB via ffmpeg
This commit is contained in:
1 parent
7932119e17
commit
6f86bca204
2 files changed
+196
-1
No files matched your search
@@ -3,7 +3,9 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from flask import (
|
||||
@@ -204,6 +206,41 @@ def jobs_list():
|
||||
return render_template("jobs.html", jobs=jobs)
|
||||
|
||||
|
||||
def _compress_for_download(file_path: str, max_bytes: int = 250 * 1024 * 1024) -> str | None:
|
||||
"""Compress file_path with ffmpeg if larger than max_bytes.
|
||||
|
||||
Returns path to `<name>_compressed.mp4` next to the input, or None to
|
||||
fall back to the original file (small file, ffmpeg missing/failed).
|
||||
"""
|
||||
if not os.path.isfile(file_path):
|
||||
return None
|
||||
if os.path.getsize(file_path) <= max_bytes:
|
||||
return None
|
||||
out_path = os.path.splitext(file_path)[0] + "_compressed.mp4"
|
||||
if os.path.isfile(out_path):
|
||||
return out_path
|
||||
try:
|
||||
probe = subprocess.run(
|
||||
["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", file_path],
|
||||
capture_output=True, text=True, timeout=60,
|
||||
)
|
||||
duration = float(json.loads(probe.stdout).get("format", {}).get("duration") or 0)
|
||||
except (OSError, subprocess.SubprocessError, ValueError, TypeError, AttributeError):
|
||||
return None
|
||||
if duration <= 0:
|
||||
return None
|
||||
bitrate = int((max_bytes * 8 * 0.92) / duration)
|
||||
try:
|
||||
subprocess.run(
|
||||
["ffmpeg", "-y", "-i", file_path, "-c:v", "libx264",
|
||||
"-b:v", str(bitrate), "-c:a", "aac", "-b:a", "64k", out_path],
|
||||
capture_output=True, timeout=3600, check=True,
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return None
|
||||
return out_path if os.path.isfile(out_path) else None
|
||||
|
||||
|
||||
@app.route("/download/<job_id>/<filename>")
|
||||
def download(job_id, filename):
|
||||
job = job_queue.get_job(job_id)
|
||||
@@ -216,7 +253,8 @@ def download(job_id, filename):
|
||||
return "File not found", 404
|
||||
if not os.path.isfile(file_path):
|
||||
return "File not found", 404
|
||||
return send_file(file_path, as_attachment=True)
|
||||
compressed = _compress_for_download(file_path)
|
||||
return send_file(compressed or file_path, as_attachment=True)
|
||||
|
||||
|
||||
@app.route("/api/videos")
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Tests for download compression to max 250MB (Task 3)."""
|
||||
|
||||
import inspect
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
import pytest
|
||||
from unittest import mock
|
||||
|
||||
import app as app_module
|
||||
from app import _compress_for_download
|
||||
|
||||
|
||||
def test_helper_exists_with_max_bytes_default():
|
||||
"""_compress_for_download accepts max_bytes, defaults to 250MB."""
|
||||
sig = inspect.signature(_compress_for_download)
|
||||
assert "max_bytes" in sig.parameters
|
||||
assert sig.parameters["max_bytes"].default == 250 * 1024 * 1024
|
||||
|
||||
|
||||
def test_small_file_returns_none(tmp_path):
|
||||
"""Files under the limit are served untouched (no compression)."""
|
||||
f = tmp_path / "small.mp4"
|
||||
f.write_bytes(b"tiny")
|
||||
assert _compress_for_download(str(f)) is None
|
||||
|
||||
|
||||
def test_file_exactly_at_limit_returns_none(tmp_path):
|
||||
"""Size == max_bytes means no compression."""
|
||||
f = tmp_path / "edge.mp4"
|
||||
f.write_bytes(b"x" * 100)
|
||||
assert _compress_for_download(str(f), max_bytes=100) is None
|
||||
|
||||
|
||||
def test_missing_file_returns_none(tmp_path):
|
||||
assert _compress_for_download(str(tmp_path / "nope.mp4")) is None
|
||||
|
||||
|
||||
def test_compression_runs_ffmpeg_with_target_bitrate(tmp_path):
|
||||
"""Oversized file triggers ffprobe + ffmpeg at computed bitrate."""
|
||||
f = tmp_path / "big.mp4"
|
||||
f.write_bytes(b"x" * 5000)
|
||||
|
||||
def fake_run(cmd, **kwargs):
|
||||
if cmd[0] == "ffprobe":
|
||||
return subprocess.CompletedProcess(
|
||||
cmd, 0, stdout='{"format": {"duration": "2.0"}}')
|
||||
open(cmd[-1], "wb").write(b"compressed")
|
||||
return subprocess.CompletedProcess(cmd, 0)
|
||||
|
||||
with mock.patch.object(app_module.subprocess, "run", side_effect=fake_run) as run:
|
||||
out = _compress_for_download(str(f), max_bytes=1024)
|
||||
|
||||
assert out == str(tmp_path / "big_compressed.mp4")
|
||||
assert os.path.isfile(out)
|
||||
cmd = run.call_args[0][0]
|
||||
assert cmd[0] == "ffmpeg"
|
||||
# bitrate = int((1024 * 8 * 0.92) / 2.0) = 3768
|
||||
assert "-b:v" in cmd
|
||||
assert cmd[cmd.index("-b:v") + 1] == "3768"
|
||||
assert cmd[cmd.index("-i") + 1] == str(f)
|
||||
assert cmd[-1] == str(tmp_path / "big_compressed.mp4")
|
||||
|
||||
|
||||
def test_existing_compressed_file_reused(tmp_path):
|
||||
"""Cache: <name>_compressed.mp4 next to input is served as-is."""
|
||||
f = tmp_path / "big.mp4"
|
||||
f.write_bytes(b"x" * 5000)
|
||||
cached = tmp_path / "big_compressed.mp4"
|
||||
cached.write_bytes(b"cached")
|
||||
|
||||
with mock.patch.object(app_module.subprocess, "run") as run:
|
||||
out = _compress_for_download(str(f), max_bytes=1024)
|
||||
|
||||
assert out == str(cached)
|
||||
run.assert_not_called()
|
||||
|
||||
|
||||
def test_ffprobe_failure_returns_none(tmp_path):
|
||||
"""ffprobe failing (bad file / missing binary) falls back to original."""
|
||||
f = tmp_path / "big.mp4"
|
||||
f.write_bytes(b"x" * 5000)
|
||||
|
||||
with mock.patch.object(
|
||||
app_module.subprocess, "run",
|
||||
side_effect=subprocess.CalledProcessError(1, "ffprobe"),
|
||||
):
|
||||
assert _compress_for_download(str(f), max_bytes=1024) is None
|
||||
|
||||
|
||||
def test_ffmpeg_missing_binary_returns_none(tmp_path):
|
||||
f = tmp_path / "big.mp4"
|
||||
f.write_bytes(b"x" * 5000)
|
||||
|
||||
with mock.patch.object(app_module.subprocess, "run", side_effect=FileNotFoundError):
|
||||
assert _compress_for_download(str(f), max_bytes=1024) is None
|
||||
|
||||
|
||||
def test_ffmpeg_failure_returns_none(tmp_path):
|
||||
f = tmp_path / "big.mp4"
|
||||
f.write_bytes(b"x" * 5000)
|
||||
|
||||
def probe_ok(cmd, **kwargs):
|
||||
if cmd[0] == "ffprobe":
|
||||
class R:
|
||||
returncode = 0
|
||||
stdout = '{"format": {"duration": "2.0"}}'
|
||||
return R()
|
||||
raise subprocess.CalledProcessError(1, cmd)
|
||||
|
||||
with mock.patch.object(app_module.subprocess, "run", side_effect=probe_ok):
|
||||
assert _compress_for_download(str(f), max_bytes=1024) is None
|
||||
|
||||
|
||||
# ── download route ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
from app import app
|
||||
app.config["TESTING"] = True
|
||||
with app.test_client() as c:
|
||||
yield c
|
||||
|
||||
|
||||
def _completed_job_with_file(name="out.mp4", content=b"video-bytes"):
|
||||
import time
|
||||
from app import job_queue
|
||||
from src.job import JobStatus
|
||||
|
||||
job = job_queue.add_job(video_path="/tmp/test.mp4", model_configs=[])
|
||||
time.sleep(0.3)
|
||||
with job_queue._lock:
|
||||
job.status = JobStatus.COMPLETED
|
||||
path = os.path.join(job.output_dir, name)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(content)
|
||||
return job, path
|
||||
|
||||
|
||||
def test_download_route_uses_compressed_file(client):
|
||||
"""Route sends the compressed path when helper returns one."""
|
||||
job, path = _completed_job_with_file()
|
||||
with mock.patch.object(app_module, "_compress_for_download", return_value=path) as m:
|
||||
resp = client.get(f"/download/{job.job_id}/out.mp4")
|
||||
assert resp.status_code == 200
|
||||
m.assert_called_once_with(path)
|
||||
assert resp.data == b"video-bytes"
|
||||
|
||||
|
||||
def test_download_route_falls_back_to_original(client):
|
||||
"""Route serves original file when helper returns None."""
|
||||
job, path = _completed_job_with_file()
|
||||
with mock.patch.object(app_module, "_compress_for_download", return_value=None):
|
||||
resp = client.get(f"/download/{job.job_id}/out.mp4")
|
||||
assert resp.status_code == 200
|
||||
assert resp.data == b"video-bytes"
|
||||
Reference in new issue
Block a user