File size: 4,402 Bytes
6cfe55f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | import importlib.util
import pathlib
import re
import sys
import tempfile
import types
import unittest
from unittest.mock import patch
def _install_stubs():
app_mod = types.ModuleType("app")
utils_pkg = types.ModuleType("app.utils")
logger_mod = types.ModuleType("app.utils.logger")
class _Logger:
@staticmethod
def info(*_args, **_kwargs):
return None
@staticmethod
def warning(*_args, **_kwargs):
return None
@staticmethod
def error(*_args, **_kwargs):
return None
def _get_logger(_name):
return _Logger()
logger_mod.get_logger = _get_logger
path_helper_mod = types.ModuleType("app.utils.path_helper")
ffmpeg_mod = types.ModuleType("ffmpeg")
pil_mod = types.ModuleType("PIL")
pil_image_mod = types.ModuleType("PIL.Image")
pil_draw_mod = types.ModuleType("PIL.ImageDraw")
pil_font_mod = types.ModuleType("PIL.ImageFont")
class _FakeImage:
pass
class _FakeImageDraw:
@staticmethod
def Draw(*_args, **_kwargs):
return None
class _FakeImageFont:
@staticmethod
def truetype(*_args, **_kwargs):
return None
@staticmethod
def load_default():
return None
pil_image_mod.Image = _FakeImage
pil_draw_mod.ImageDraw = _FakeImageDraw
pil_font_mod.ImageFont = _FakeImageFont
def _get_app_dir(name):
return name
path_helper_mod.get_app_dir = _get_app_dir
ffmpeg_mod.probe = lambda *_args, **_kwargs: {"format": {"duration": "0"}}
sys.modules.setdefault("app", app_mod)
sys.modules.setdefault("app.utils", utils_pkg)
sys.modules["PIL"] = pil_mod
sys.modules["PIL.Image"] = pil_image_mod
sys.modules["PIL.ImageDraw"] = pil_draw_mod
sys.modules["PIL.ImageFont"] = pil_font_mod
sys.modules["ffmpeg"] = ffmpeg_mod
sys.modules["app.utils.logger"] = logger_mod
sys.modules["app.utils.path_helper"] = path_helper_mod
def _load_video_reader_module():
_install_stubs()
root = pathlib.Path(__file__).resolve().parents[1]
module_path = root / "app" / "utils" / "video_reader.py"
spec = importlib.util.spec_from_file_location("video_reader", module_path)
if spec is None or spec.loader is None:
raise ImportError("video_reader module spec not found")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
video_reader_module = _load_video_reader_module()
VideoReader = video_reader_module.VideoReader
def _make_fake_ffmpeg_runner(colors_by_second):
def _runner(cmd, check=True):
output_path = next((arg for arg in cmd if isinstance(arg, str) and arg.endswith(".jpg")), None)
if output_path is None:
raise AssertionError("Output path not found in ffmpeg cmd")
match = re.search(r"frame_(\d{2})_(\d{2})\.jpg$", output_path)
if match is None:
raise AssertionError("Unexpected output path")
sec = int(match.group(1)) * 60 + int(match.group(2))
payload = colors_by_second[sec]
with open(output_path, "wb") as f:
f.write(payload)
return 0
return _runner
class TestVideoReaderDeduplicateFrames(unittest.TestCase):
def test_extract_frames_skips_adjacent_duplicates_when_enabled(self):
with tempfile.TemporaryDirectory() as tmp_dir:
frame_dir = pathlib.Path(tmp_dir) / "frames"
grid_dir = pathlib.Path(tmp_dir) / "grids"
reader = VideoReader(
video_path="dummy.mp4",
frame_interval=1,
frame_dir=str(frame_dir),
grid_dir=str(grid_dir),
)
fake_colors = {
0: b"frame-a",
1: b"frame-a",
2: b"frame-b",
3: b"frame-b",
}
with patch.object(video_reader_module.ffmpeg, "probe", return_value={"format": {"duration": "4"}}), \
patch.object(video_reader_module.subprocess, "run", side_effect=_make_fake_ffmpeg_runner(fake_colors)):
paths = reader.extract_frames(max_frames=10)
names = [pathlib.Path(p).name for p in paths]
self.assertEqual(names, ["frame_00_00.jpg", "frame_00_02.jpg"])
if __name__ == "__main__":
unittest.main()
|