fix: atomic compression writes, enforce size limit, job status guard

tests/test_download_compress.py

@@ -2,7 +2,10 @@
 
 import inspect
 import os
+import shutil
 import subprocess
+import threading
+import time
 
 import pytest
 from unittest import mock
@@ -10,111 +13,254 @@ 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."""
+    """Oversized file triggers ffmpeg at video bitrate net of audio budget."""
     f = tmp_path / "big.mp4"
-    f.write_bytes(b"x" * 5000)
+    f.write_bytes(b"x" * 300_000)
 
     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)
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         mock.patch.object(app_module.subprocess, "run", side_effect=fake_run) as run:
+        out = _compress_for_download(str(f), max_bytes=200_000)
 
     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
+    # bitrate = int((200_000 * 8 * 0.92) / 2.0) - 64_000 = 672000
     assert "-b:v" in cmd
-    assert cmd[cmd.index("-b:v") + 1] == "3768"
+    assert cmd[cmd.index("-b:v") + 1] == "672000"
     assert cmd[cmd.index("-i") + 1] == str(f)
-    assert cmd[-1] == str(tmp_path / "big_compressed.mp4")
+    # audio stays budgeted
+    assert cmd[cmd.index("-b:a") + 1] == "64000"
+    # Jetson: fast preset, explicit container (output is a .tmp path)
+    assert cmd[cmd.index("-preset") + 1] == "veryfast"
+    assert cmd[cmd.index("-f") + 1] == "mp4"
+    # encode goes to a temp path, never the final path directly
+    assert cmd[-1] == str(tmp_path / "big_compressed.mp4.tmp")
+    assert run.call_args[1]["stdin"] is subprocess.DEVNULL
 
+def test_minimum_video_bitrate_floor(tmp_path):
+    """Tiny budget still gets a usable floor bitrate."""
+    f = tmp_path / "big.mp4"
+    f.write_bytes(b"x" * 5000)
+
+    def fake_run(cmd, **kwargs):
+        open(cmd[-1], "wb").write(b"c")
+        return subprocess.CompletedProcess(cmd, 0)
+
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 100.0}), \
+         mock.patch.object(app_module.subprocess, "run", side_effect=fake_run) as run:
+        _compress_for_download(str(f), max_bytes=1024)
+
+    cmd = run.call_args[0][0]
+    assert int(cmd[cmd.index("-b:v") + 1]) == 100_000
+
+def test_partial_encode_never_cached(tmp_path):
+    """ffmpeg writes only to .tmp; final path appears only after success."""
+    f = tmp_path / "big.mp4"
+    f.write_bytes(b"x" * 300_000)
+    final = str(tmp_path / "big_compressed.mp4")
+    seen = {}
+
+    def fake_run(cmd, **kwargs):
+        seen["final_during_encode"] = os.path.isfile(final)
+        seen["tmp"] = cmd[-1]
+        open(cmd[-1], "wb").write(b"compressed")
+        return subprocess.CompletedProcess(cmd, 0)
+
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         mock.patch.object(app_module.subprocess, "run", side_effect=fake_run):
+        out = _compress_for_download(str(f), max_bytes=200_000)
+
+    assert out == final
+    assert seen["final_during_encode"] is False
+    assert seen["tmp"] == final + ".tmp"
+    assert os.path.isfile(final)
+    assert not os.path.exists(final + ".tmp")
+
+def test_failed_encode_leaves_no_partial_file(tmp_path):
+    """Interrupted/killed encode must not leave a servable artifact."""
+    f = tmp_path / "big.mp4"
+    f.write_bytes(b"x" * 5000)
+    final = str(tmp_path / "big_compressed.mp4")
+
+    def fake_run(cmd, **kwargs):
+        open(cmd[-1], "wb").write(b"partial")
+        raise subprocess.CalledProcessError(1, cmd)
+
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         mock.patch.object(app_module.subprocess, "run", side_effect=fake_run):
+        assert _compress_for_download(str(f), max_bytes=200_000) is None
+
+    assert not os.path.exists(final)
+    assert not os.path.exists(final + ".tmp")
+
+def test_oversized_result_discarded(tmp_path):
+    """If the encode still exceeds max_bytes, drop it and fall back."""
+    f = tmp_path / "big.mp4"
+    f.write_bytes(b"x" * 5000)
+    final = str(tmp_path / "big_compressed.mp4")
+
+    def fake_run(cmd, **kwargs):
+        open(cmd[-1], "wb").write(b"y" * 4096)
+        return subprocess.CompletedProcess(cmd, 0)
+
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         mock.patch.object(app_module.subprocess, "run", side_effect=fake_run):
+        assert _compress_for_download(str(f), max_bytes=1024) is None
+
+    assert not os.path.exists(final)
+    assert not os.path.exists(final + ".tmp")
 
 def test_existing_compressed_file_reused(tmp_path):
-    """Cache: <name>_compressed.mp4 next to input is served as-is."""
+    """Cache: fresh <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")
+    os.utime(cached, (f.stat().st_mtime + 10, f.stat().st_mtime + 10))
 
-    with mock.patch.object(app_module.subprocess, "run") as run:
+    with mock.patch.object(app_module, "probe_video") as probe, \
+         mock.patch.object(app_module.subprocess, "run") as run:
         out = _compress_for_download(str(f), max_bytes=1024)
 
     assert out == str(cached)
+    probe.assert_not_called()
     run.assert_not_called()
 
+def test_stale_cache_reencoded(tmp_path):
+    """Cache older than the source video is invalid and gets re-encoded."""
+    f = tmp_path / "big.mp4"
+    f.write_bytes(b"x" * 300_000)
+    cached = tmp_path / "big_compressed.mp4"
+    cached.write_bytes(b"old")
+    os.utime(cached, (1000, 1000))  # far older than source
 
-def test_ffprobe_failure_returns_none(tmp_path):
-    """ffprobe failing (bad file / missing binary) falls back to original."""
+    def fake_run(cmd, **kwargs):
+        open(cmd[-1], "wb").write(b"fresh")
+        return subprocess.CompletedProcess(cmd, 0)
+
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         mock.patch.object(app_module.subprocess, "run", side_effect=fake_run) as run:
+        out = _compress_for_download(str(f), max_bytes=200_000)
+
+    assert out == str(cached)
+    run.assert_called_once()
+    assert cached.read_bytes() == b"fresh"
+
+def test_probe_failure_returns_none(tmp_path):
+    """probe failure falls back to the original file (no compression)."""
     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"),
-    ):
+    with mock.patch.object(app_module, "probe_video", side_effect=FileNotFoundError):
         assert _compress_for_download(str(f), max_bytes=1024) is None
 
+def test_nonpositive_duration_returns_none(tmp_path):
+    """duration <= 0 means no compression (nothing to budget against)."""
+    f = tmp_path / "big.mp4"
+    f.write_bytes(b"x" * 5000)
+
+    for duration in (0, 0.0, -1.0, None):
+        with mock.patch.object(app_module, "probe_video", return_value={"duration": duration}):
+            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):
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         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):
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         mock.patch.object(app_module.subprocess, "run",
+                           side_effect=subprocess.CalledProcessError(1, "ffmpeg")):
         assert _compress_for_download(str(f), max_bytes=1024) is None
 
+def test_concurrent_downloads_encode_once(tmp_path):
+    """Threaded requests for the same file share one encode (global lock)."""
+    f = tmp_path / "big.mp4"
+    f.write_bytes(b"x" * 300_000)
+    calls = []
+
+    def slow_run(cmd, **kwargs):
+        calls.append(cmd)
+        time.sleep(0.3)
+        open(cmd[-1], "wb").write(b"compressed")
+        return subprocess.CompletedProcess(cmd, 0)
+
+    results = []
+
+    def worker():
+        results.append(_compress_for_download(str(f), max_bytes=200_000))
+
+    with mock.patch.object(app_module, "probe_video", return_value={"duration": 2.0}), \
+         mock.patch.object(app_module.subprocess, "run", side_effect=slow_run):
+        threads = [threading.Thread(target=worker) for _ in range(2)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join(timeout=30)
+
+    assert len(calls) == 1
+    assert results == [str(tmp_path / "big_compressed.mp4")] * 2
+
+# ── real ffmpeg smoke test ─────────────────────────────────────────────
+
+@pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg not installed")
+def test_real_ffmpeg_smoke(tmp_path):
+    """Real encode of a tiny video+audio file lands under the byte cap."""
+    src = tmp_path / "tiny.mp4"
+    # lossless source so it is genuinely larger than max_bytes
+    subprocess.run(
+        ["ffmpeg", "-y",
+         "-f", "lavfi", "-i", "testsrc=duration=2:size=320x240:rate=10",
+         "-f", "lavfi", "-i", "sine=frequency=440:duration=2",
+         "-c:v", "libx264", "-qp", "0", "-preset", "veryfast",
+         "-c:a", "aac", "-b:a", "128k", "-shortest", str(src)],
+        check=True, capture_output=True, stdin=subprocess.DEVNULL,
+    )
+    max_bytes = 60_000
+    assert os.path.getsize(src) > max_bytes, "source must exceed cap to compress"
+
+    out = _compress_for_download(str(src), max_bytes=max_bytes)
+
+    assert out is not None
+    assert os.path.getsize(out) <= max_bytes
+    assert not os.path.exists(out + ".tmp")
 
 # ── download route ──────────────────────────────────────────────────────
 
-
 @pytest.fixture
 def client():
     from app import app
@@ -122,21 +268,21 @@ def client():
     with app.test_client() as c:
         yield c
 
-
-def _completed_job_with_file(name="out.mp4", content=b"video-bytes"):
-    import time
+def _job_with_file(status="COMPLETED", name="out.mp4", content=b"video-bytes"):
     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
+        job.status = JobStatus[status]
     path = os.path.join(job.output_dir, name)
     with open(path, "wb") as fh:
         fh.write(content)
     return job, path
 
+def _completed_job_with_file(name="out.mp4", content=b"video-bytes"):
+    return _job_with_file("COMPLETED", name, content)
 
 def test_download_route_uses_compressed_file(client):
     """Route sends the compressed path when helper returns one."""
@@ -147,7 +293,6 @@ def test_download_route_uses_compressed_file(client):
     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()
@@ -155,3 +300,16 @@ def test_download_route_falls_back_to_original(client):
         resp = client.get(f"/download/{job.job_id}/out.mp4")
     assert resp.status_code == 200
     assert resp.data == b"video-bytes"
+
+def test_download_route_409_when_job_running(client):
+    """No downloads while the job is still writing its output."""
+    job, _path = _job_with_file("RUNNING")
+    with mock.patch.object(app_module, "_compress_for_download") as m:
+        resp = client.get(f"/download/{job.job_id}/out.mp4")
+    assert resp.status_code == 409
+    m.assert_not_called()
+
+def test_download_route_409_when_job_pending(client):
+    job, _path = _job_with_file("PENDING")
+    resp = client.get(f"/download/{job.job_id}/out.mp4")
+    assert resp.status_code == 409