From 6f86bca204643b701f671197b285aab51c6b8960 Mon Sep 17 00:00:00 2001 From: jetson Date: Mon, 28 Sep 2026 09:39:56 +0700 Subject: [PATCH] feat: compress downloads to max 250MB via ffmpeg --- app.py | 40 +++++++- tests/test_download_compress.py | 157 ++++++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+), 1 deletion(-) create mode 100644 tests/test_download_compress.py diff --git a/app.py b/app.py index cfb4335..5217bb5 100644 --- a/app.py +++ b/app.py @@ -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 `_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//") 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") diff --git a/tests/test_download_compress.py b/tests/test_download_compress.py new file mode 100644 index 0000000..2004fe4 --- /dev/null +++ b/tests/test_download_compress.py @@ -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: _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"