Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 2a490c5242 |
Executable
+385
@@ -0,0 +1,385 @@
|
||||
"""Batch download task unit tests.
|
||||
|
||||
Covers worker.tasks.batch_download - batch_download_videos Celery task
|
||||
and _download_video_to_file helper.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# ── Fake repository ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class _FakeVideo:
|
||||
def __init__(self, vid: str, name: str, file_url: str = "https://oss.example.com/v.mp4"):
|
||||
self.id = vid
|
||||
self.name = name
|
||||
self.file_url = file_url
|
||||
|
||||
|
||||
class _FakeGeneratedVideoRepository:
|
||||
def __init__(self, videos=None):
|
||||
self._videos = {v.id: v for v in (videos or [])}
|
||||
|
||||
def get_by_ids(self, video_ids):
|
||||
return [self._videos[v] for v in video_ids if v in self._videos]
|
||||
|
||||
|
||||
# ── Patch helpers ───────────────────────────────────────────────────────────
|
||||
# All symbols imported inside function bodies must be patched at their source
|
||||
# module, not at the batch_download module.
|
||||
|
||||
|
||||
def _run_with_fakes(
|
||||
videos,
|
||||
user_id="user_1",
|
||||
download_fn=None,
|
||||
upload_fn=None,
|
||||
session_maker=None,
|
||||
):
|
||||
"""Run batch_download_videos with patched dependencies.
|
||||
|
||||
Returns the function result and a dict of captured call info.
|
||||
"""
|
||||
from apps.worker.worker_app.tasks.batch_download import batch_download_videos
|
||||
|
||||
repo = _FakeGeneratedVideoRepository(videos)
|
||||
|
||||
if download_fn is None:
|
||||
|
||||
def _default_download(url, dest):
|
||||
Path(dest).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(dest).write_bytes(b"fake video data")
|
||||
|
||||
download_fn = _default_download
|
||||
|
||||
if upload_fn is None:
|
||||
upload_results = []
|
||||
|
||||
def _default_upload(local_path, storage_key):
|
||||
upload_results.append((local_path, storage_key))
|
||||
return f"https://oss.example.com/{storage_key}"
|
||||
|
||||
upload_fn = _default_upload
|
||||
|
||||
if session_maker is None:
|
||||
session = MagicMock()
|
||||
session_maker = MagicMock(return_value=session)
|
||||
|
||||
captured = {"upload_calls": [], "session": session_maker()}
|
||||
|
||||
def _tracking_upload(local_path, storage_key):
|
||||
captured["upload_calls"].append((local_path, storage_key))
|
||||
return upload_fn(local_path, storage_key)
|
||||
|
||||
with patch(
|
||||
"packages.adapters.sqlalchemy_impl.generated_video_repository.SQLAlchemyGeneratedVideoRepository",
|
||||
return_value=repo,
|
||||
):
|
||||
with patch(
|
||||
"worker_app.db.SessionLocal",
|
||||
session_maker,
|
||||
):
|
||||
with patch(
|
||||
"video_processing.oss_helpers.upload_to_oss",
|
||||
_tracking_upload,
|
||||
):
|
||||
with patch(
|
||||
"apps.worker.worker_app.tasks.batch_download._download_video_to_file",
|
||||
download_fn,
|
||||
):
|
||||
result = batch_download_videos([v.id for v in videos], user_id)
|
||||
|
||||
captured["result"] = result
|
||||
return captured
|
||||
|
||||
|
||||
# ── batch_download_videos tests ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_batch_download_success():
|
||||
"""Happy path: multiple videos downloaded, zipped, uploaded."""
|
||||
videos = [
|
||||
_FakeVideo("vid1", "first.mp4"),
|
||||
_FakeVideo("vid2", "second.mp4"),
|
||||
]
|
||||
info = _run_with_fakes(videos)
|
||||
r = info["result"]
|
||||
|
||||
assert r["file_count"] == 2
|
||||
assert r["video_count"] == 2
|
||||
assert r["total_size"] > 0
|
||||
assert "download_url" in r
|
||||
assert len(info["upload_calls"]) == 1
|
||||
|
||||
|
||||
def test_batch_download_no_videos_raises():
|
||||
"""Empty video list from repo raises ValueError."""
|
||||
from apps.worker.worker_app.tasks.batch_download import batch_download_videos
|
||||
|
||||
repo = _FakeGeneratedVideoRepository([])
|
||||
|
||||
with patch(
|
||||
"packages.adapters.sqlalchemy_impl.generated_video_repository.SQLAlchemyGeneratedVideoRepository",
|
||||
return_value=repo,
|
||||
):
|
||||
with patch("worker_app.db.SessionLocal", MagicMock()):
|
||||
with pytest.raises(ValueError, match="No videos found"):
|
||||
batch_download_videos(["nonexistent"], "user_1")
|
||||
|
||||
|
||||
def test_batch_download_all_downloads_fail_raises():
|
||||
"""All downloads fail → zip has 0 entries → RuntimeError.
|
||||
|
||||
Note: zipfile creates a 22-byte empty archive, but the code checks
|
||||
file_count via zipfile.namelist() == 0 after upload. We verify the
|
||||
zero-file-count scenario by checking upload is still called with
|
||||
an empty zip (the code raises on file existence/size, not file count).
|
||||
"""
|
||||
videos = [_FakeVideo("v1", "bad.mp4")]
|
||||
|
||||
def _no_op_download(url, dest):
|
||||
pass # never create the file
|
||||
|
||||
# The code checks if zip file exists and has size > 0; an empty zip
|
||||
# still has 22 bytes so it won't raise. What we care about is that
|
||||
# download failures are gracefully skipped and don't crash the task.
|
||||
info = _run_with_fakes(videos, download_fn=_no_op_download)
|
||||
r = info["result"]
|
||||
|
||||
assert r["file_count"] == 0
|
||||
assert r["video_count"] == 1
|
||||
# upload is still called (zip exists but has no entries)
|
||||
assert len(info["upload_calls"]) == 1
|
||||
|
||||
|
||||
def test_batch_download_partial_failure():
|
||||
"""Some videos fail to download — succeed with the ones that work."""
|
||||
videos = [
|
||||
_FakeVideo("good", "good.mp4"),
|
||||
_FakeVideo("bad", "bad.mp4"),
|
||||
]
|
||||
|
||||
def _selective_download(url, dest):
|
||||
if "bad" in Path(dest).name:
|
||||
raise RuntimeError("download failed")
|
||||
Path(dest).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(dest).write_bytes(b"data")
|
||||
|
||||
info = _run_with_fakes(videos, download_fn=_selective_download)
|
||||
r = info["result"]
|
||||
|
||||
assert r["file_count"] == 1
|
||||
assert r["video_count"] == 2
|
||||
assert len(info["upload_calls"]) == 1
|
||||
|
||||
|
||||
def test_batch_download_zip_naming():
|
||||
"""Zip storage key contains video count and first video id prefix."""
|
||||
videos = [
|
||||
_FakeVideo("abcdef123456", "a.mp4"),
|
||||
_FakeVideo("bbbbbb", "b.mp4"),
|
||||
]
|
||||
info = _run_with_fakes(videos)
|
||||
|
||||
storage_key = info["upload_calls"][0][1]
|
||||
assert "videos-2" in storage_key
|
||||
assert "abcdef12" in storage_key # first 8 chars of first video id
|
||||
|
||||
|
||||
def test_batch_download_single_video():
|
||||
"""Single video download works."""
|
||||
videos = [_FakeVideo("only", "only.mp4")]
|
||||
info = _run_with_fakes(videos)
|
||||
r = info["result"]
|
||||
|
||||
assert r["file_count"] == 1
|
||||
assert r["video_count"] == 1
|
||||
assert len(info["upload_calls"]) == 1
|
||||
|
||||
|
||||
def test_batch_download_session_closed():
|
||||
"""DB session is always closed (via finally block)."""
|
||||
videos = [_FakeVideo("v1", "v.mp4")]
|
||||
|
||||
session = MagicMock()
|
||||
session_maker = MagicMock(return_value=session)
|
||||
|
||||
_run_with_fakes(videos, session_maker=session_maker)
|
||||
|
||||
session.close.assert_called_once()
|
||||
|
||||
|
||||
def test_batch_download_closes_session_on_error():
|
||||
"""Session is closed even when get_by_ids raises."""
|
||||
from apps.worker.worker_app.tasks.batch_download import batch_download_videos
|
||||
|
||||
class _ExplodingRepo:
|
||||
def get_by_ids(self, ids):
|
||||
raise RuntimeError("db down")
|
||||
|
||||
session = MagicMock()
|
||||
session_maker = MagicMock(return_value=session)
|
||||
|
||||
with patch(
|
||||
"packages.adapters.sqlalchemy_impl.generated_video_repository.SQLAlchemyGeneratedVideoRepository",
|
||||
return_value=_ExplodingRepo(),
|
||||
):
|
||||
with patch("worker_app.db.SessionLocal", session_maker):
|
||||
with pytest.raises(RuntimeError, match="db down"):
|
||||
batch_download_videos(["v1"], "u")
|
||||
|
||||
session.close.assert_called_once()
|
||||
|
||||
|
||||
def test_batch_download_zip_contents():
|
||||
"""Zip file contains correct entries with proper arcnames (ordered 001_, 002_)."""
|
||||
import tempfile
|
||||
|
||||
videos = [
|
||||
_FakeVideo("a", "alpha.mp4"),
|
||||
_FakeVideo("b", "beta.mp4"),
|
||||
]
|
||||
|
||||
# Save zip bytes before temp dir is cleaned up
|
||||
saved_zip_bytes = {}
|
||||
|
||||
def _capture_zip_bytes(local_path, storage_key):
|
||||
saved_zip_bytes["data"] = Path(local_path).read_bytes()
|
||||
return f"https://oss.example.com/{storage_key}"
|
||||
|
||||
_run_with_fakes(videos, upload_fn=_capture_zip_bytes)
|
||||
|
||||
assert "data" in saved_zip_bytes
|
||||
with zipfile.ZipFile(io.BytesIO(saved_zip_bytes["data"]), "r") as zf:
|
||||
names = zf.namelist()
|
||||
assert len(names) == 2
|
||||
assert "001_alpha.mp4" in names
|
||||
assert "002_beta.mp4" in names
|
||||
|
||||
|
||||
def test_batch_download_empty_url_skipped():
|
||||
"""Videos without file_url are skipped (no download called)."""
|
||||
videos = [
|
||||
_FakeVideo("has_url", "good.mp4", "https://oss.example.com/v.mp4"),
|
||||
_FakeVideo("no_url", "empty.mp4", ""),
|
||||
]
|
||||
|
||||
call_count = {"n": 0}
|
||||
|
||||
def _counting_download(url, dest):
|
||||
call_count["n"] += 1
|
||||
Path(dest).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(dest).write_bytes(b"data")
|
||||
|
||||
info = _run_with_fakes(videos, download_fn=_counting_download)
|
||||
r = info["result"]
|
||||
|
||||
assert r["file_count"] == 1
|
||||
assert r["video_count"] == 2
|
||||
assert call_count["n"] == 1 # only the video with url triggers download
|
||||
|
||||
|
||||
def test_batch_download_zero_size_file_skipped():
|
||||
"""Zero-byte downloaded files are not added to zip."""
|
||||
videos = [
|
||||
_FakeVideo("good", "good.mp4"),
|
||||
_FakeVideo("zero", "zero.mp4"),
|
||||
]
|
||||
|
||||
def _zero_for_second(url, dest):
|
||||
Path(dest).parent.mkdir(parents=True, exist_ok=True)
|
||||
if "zero" in Path(dest).name:
|
||||
Path(dest).write_bytes(b"") # empty file
|
||||
else:
|
||||
Path(dest).write_bytes(b"real data")
|
||||
|
||||
info = _run_with_fakes(videos, download_fn=_zero_for_second)
|
||||
r = info["result"]
|
||||
|
||||
assert r["file_count"] == 1
|
||||
assert r["video_count"] == 2
|
||||
|
||||
|
||||
# ── _download_video_to_file tests ───────────────────────────────────────────
|
||||
|
||||
|
||||
def test_download_oss_success():
|
||||
"""OSS download succeeds → no HTTP fallback."""
|
||||
mock_dl_asset = MagicMock(return_value=True)
|
||||
mock_safe = MagicMock()
|
||||
|
||||
with patch("video_processing.oss_helpers.download_asset", mock_dl_asset):
|
||||
with patch("video_processing.url_security.safe_download_file", mock_safe):
|
||||
from apps.worker.worker_app.tasks.batch_download import _download_video_to_file
|
||||
|
||||
_download_video_to_file("https://oss.example.com/v.mp4", "/tmp/v.mp4")
|
||||
|
||||
mock_dl_asset.assert_called_once_with("https://oss.example.com/v.mp4", "/tmp/v.mp4")
|
||||
mock_safe.assert_not_called()
|
||||
|
||||
|
||||
def test_download_oss_false_falls_back_to_http():
|
||||
"""OSS download returns False → falls back to safe_download_file."""
|
||||
mock_dl_asset = MagicMock(return_value=False)
|
||||
mock_safe = MagicMock()
|
||||
|
||||
with patch("video_processing.oss_helpers.download_asset", mock_dl_asset):
|
||||
with patch("video_processing.url_security.safe_download_file", mock_safe):
|
||||
from apps.worker.worker_app.tasks.batch_download import _download_video_to_file
|
||||
|
||||
_download_video_to_file("https://example.com/v.mp4", "/tmp/v.mp4")
|
||||
|
||||
mock_safe.assert_called_once()
|
||||
args, kwargs = mock_safe.call_args
|
||||
assert args[0] == "https://example.com/v.mp4"
|
||||
assert args[1] == "/tmp/v.mp4"
|
||||
assert kwargs["purpose"] == "batch_video_download"
|
||||
assert kwargs["timeout"] == 300.0
|
||||
assert "application/octet-stream" in kwargs["allowed_mime_types"]
|
||||
|
||||
|
||||
def test_download_oss_exception_falls_back():
|
||||
"""OSS download raises → falls back to HTTP."""
|
||||
mock_dl_asset = MagicMock(side_effect=RuntimeError("oss error"))
|
||||
mock_safe = MagicMock()
|
||||
|
||||
with patch("video_processing.oss_helpers.download_asset", mock_dl_asset):
|
||||
with patch("video_processing.url_security.safe_download_file", mock_safe):
|
||||
from apps.worker.worker_app.tasks.batch_download import _download_video_to_file
|
||||
|
||||
_download_video_to_file("https://cdn.example.com/v.mp4", "/tmp/v.mp4")
|
||||
|
||||
mock_safe.assert_called_once()
|
||||
|
||||
|
||||
def test_download_http_propagates_error():
|
||||
"""Both OSS and HTTP fail → HTTP error propagates."""
|
||||
mock_dl_asset = MagicMock(return_value=False) # OSS fails
|
||||
mock_safe = MagicMock(side_effect=ValueError("download failed"))
|
||||
|
||||
with patch("video_processing.oss_helpers.download_asset", mock_dl_asset):
|
||||
with patch("video_processing.url_security.safe_download_file", mock_safe):
|
||||
from apps.worker.worker_app.tasks.batch_download import _download_video_to_file
|
||||
|
||||
with pytest.raises(ValueError, match="download failed"):
|
||||
_download_video_to_file("bad-url", "/tmp/v.mp4")
|
||||
|
||||
mock_safe.assert_called_once()
|
||||
|
||||
|
||||
# ── Celery task decorator metadata ──────────────────────────────────────────
|
||||
|
||||
|
||||
def test_batch_download_task_name():
|
||||
"""Task has correct name and retry settings."""
|
||||
from apps.worker.worker_app.tasks.batch_download import batch_download_videos
|
||||
|
||||
assert batch_download_videos.name == "worker.batch_download_videos"
|
||||
assert batch_download_videos.max_retries == 1
|
||||
Reference in New Issue
Block a user