|
1
|
"""Tests for the frame extractor module.""" |
|
2
|
|
|
3
|
import os |
|
4
|
import tempfile |
|
5
|
|
|
6
|
import numpy as np |
|
7
|
import pytest |
|
8
|
|
|
9
|
from video_processor.extractors.frame_extractor import ( |
|
10
|
calculate_frame_difference, |
|
11
|
is_gpu_available, |
|
12
|
save_frames, |
|
13
|
) |
|
14
|
|
|
15
|
|
|
16
|
# Create dummy test frames |
|
17
|
@pytest.fixture |
|
18
|
def dummy_frames(): |
|
19
|
# Create a list of dummy frames with different content |
|
20
|
frames = [] |
|
21
|
for i in range(3): |
|
22
|
# Create frame with different intensity for each |
|
23
|
frame = np.ones((100, 100, 3), dtype=np.uint8) * (i * 50) |
|
24
|
frames.append(frame) |
|
25
|
return frames |
|
26
|
|
|
27
|
|
|
28
|
def test_calculate_frame_difference(): |
|
29
|
"""Test frame difference calculation.""" |
|
30
|
# Create two frames with some difference |
|
31
|
frame1 = np.zeros((100, 100, 3), dtype=np.uint8) |
|
32
|
frame2 = np.ones((100, 100, 3), dtype=np.uint8) * 128 # 50% intensity |
|
33
|
|
|
34
|
# Calculate difference |
|
35
|
diff = calculate_frame_difference(frame1, frame2) |
|
36
|
|
|
37
|
# Expected difference is around 128/255 = 0.5 |
|
38
|
assert 0.45 <= diff <= 0.55 |
|
39
|
|
|
40
|
# Test identical frames |
|
41
|
diff_identical = calculate_frame_difference(frame1, frame1.copy()) |
|
42
|
assert diff_identical < 0.001 # Should be very close to 0 |
|
43
|
|
|
44
|
|
|
45
|
def test_is_gpu_available(): |
|
46
|
"""Test GPU availability check.""" |
|
47
|
# This just tests that the function runs without error |
|
48
|
# We don't assert the result because it depends on the system |
|
49
|
result = is_gpu_available() |
|
50
|
assert isinstance(result, bool) |
|
51
|
|
|
52
|
|
|
53
|
def test_save_frames(dummy_frames): |
|
54
|
"""Test saving frames to disk.""" |
|
55
|
with tempfile.TemporaryDirectory() as temp_dir: |
|
56
|
# Save frames |
|
57
|
paths = save_frames(dummy_frames, temp_dir, "test_frame") |
|
58
|
|
|
59
|
# Check that we got the correct number of paths |
|
60
|
assert len(paths) == len(dummy_frames) |
|
61
|
|
|
62
|
# Check that files were created |
|
63
|
for path in paths: |
|
64
|
assert os.path.exists(path) |
|
65
|
assert os.path.getsize(path) > 0 # Files should have content |
|
66
|
|