1"""Tests for `chives.media`."""
2
3from pathlib import Path
4from typing import Any
5
6from PIL import Image
7import pytest
8
9from chives.media import (
10    convert_srt_to_vtt,
11    create_image_entity,
12    create_video_entity,
13    get_media_paths,
14    is_av1_video,
15)
16
17
18@pytest.fixture
19def fixtures_dir() -> Path:
20    """
21    Return the directory where media fixtures are stored.
22    """
23    return Path("tests/fixtures/media")
24
25
26def test_is_av1_video(fixtures_dir: Path) -> None:
27    """is_av1_video correctly detects AV1 videos."""
28    # These two videos were downloaded from
29    # https://test-videos.co.uk/sintel/mp4-h264 and
30    # https://test-videos.co.uk/sintel/mp4-av1
31    assert not is_av1_video(fixtures_dir / "Sintel_360_10s_1MB_H264.mp4")
32    assert is_av1_video(fixtures_dir / "Sintel_360_10s_1MB_AV1.mp4")
33
34
35class TestCreateImageEntity:
36    """
37    Tests for create_image_entity().
38    """
39
40    def test_basic_image(self, fixtures_dir: Path) -> None:
41        """
42        Get an image entity for a basic blue square.
43        """
44        entity = create_image_entity(fixtures_dir / "blue.png")
45        assert entity == {
46            "type": "image",
47            "path": "tests/fixtures/media/blue.png",
48            "width": 32,
49            "height": 16,
50            "tint_colour": "#0000ff",
51        }
52
53    @pytest.mark.parametrize(
54        "filename",
55        [
56            # This is a solid blue image with a section in the middle deleted
57            "blue_with_hole.png",
58            #
59            # An asteroid belt drawn in TikZ by TeX.SE user Qrrbrbirlbel,
60            # which has `transparency` in its im.info.
61            # Downloaded from http://tex.stackexchange.com/a/111974/9668
62            "asteroid_belt.png",
63        ],
64    )
65    def test_image_with_transparency(self, fixtures_dir: Path, filename: str) -> None:
66        """
67        If an image has transparent pixels, then the entity has
68        `has_transparency=True`.
69        """
70        entity = create_image_entity(fixtures_dir / filename)
71        assert entity["has_transparency"]
72
73    @pytest.mark.parametrize(
74        "filename",
75        [
76            "blue.png",
77            "space.jpg",
78            #
79            # An animated electric field drawn in TikZ.
80            # Downloaded from https://tex.stackexchange.com/a/158930/9668
81            "electric_field.gif",
82        ],
83    )
84    def test_image_without_transparency(
85        self, fixtures_dir: Path, filename: str
86    ) -> None:
87        """
88        If an image has no transparent pixels, then the entity doesn't
89        have a `has_transparency` key.
90        """
91        entity = create_image_entity(fixtures_dir / filename)
92        assert "has_transparency" not in entity
93
94    # These test files were downloaded from Dave Perrett repo:
95    # https://github.com/recurser/exif-orientation-examples
96
97    @pytest.mark.parametrize(
98        "filename",
99        [
100            "Landscape_0.jpg",
101            "Landscape_1.jpg",
102            "Landscape_2.jpg",
103            "Landscape_3.jpg",
104            "Landscape_4.jpg",
105            "Landscape_5.jpg",
106            "Landscape_6.jpg",
107            "Landscape_7.jpg",
108            "Landscape_8.jpg",
109        ],
110    )
111    def test_accounts_for_exif_orientation(
112        self, fixtures_dir: Path, filename: str
113    ) -> None:
114        """
115        The dimensions are the display dimensions, which accounts for
116        the EXIF orientation.
117        """
118        entity = create_image_entity(fixtures_dir / filename)
119        assert (entity["width"], entity["height"]) == (1800, 1200)
120
121    def test_animated_image(self, fixtures_dir: Path) -> None:
122        """
123        If an image is animated, the entity has `is_animated=True`.
124        """
125        # An animated electric field drawn in TikZ.
126        # Downloaded from https://tex.stackexchange.com/a/158930/9668
127        entity = create_image_entity(fixtures_dir / "electric_field.gif")
128        assert entity["is_animated"]
129
130    def test_other_attrs_are_forwarded(self, fixtures_dir: Path) -> None:
131        """
132        The `alt_text` and `source_url` values are forwarded to the
133        final entity.
134        """
135        entity = create_image_entity(
136            fixtures_dir / "blue.png",
137            alt_text="This is the alt text",
138            source_url="https://example.com/blue.png",
139        )
140
141        assert entity["alt_text"] == "This is the alt text"
142        assert entity["source_url"] == "https://example.com/blue.png"
143
144    def test_alt_text_and_generate_transcript_is_error(
145        self, fixtures_dir: Path
146    ) -> None:
147        """
148        You can't pass `alt_text` and `generate_transcript` at the same time.
149        """
150        with pytest.raises(TypeError):
151            create_image_entity(
152                fixtures_dir / "blue.png",
153                alt_text="This is the alt text",
154                generate_transcript=True,
155            )
156
157    def test_generate_transcript(self, fixtures_dir: Path) -> None:
158        """
159        If you pass `generate_transcript=True`, the image is OCR'd for alt text.
160        """
161        entity = create_image_entity(
162            fixtures_dir / "underlined_text.png", generate_transcript=True
163        )
164        assert entity["alt_text"] == "I visited Berlin in Germany."
165
166    def test_generate_transcript_if_no_text(self, fixtures_dir: Path) -> None:
167        """
168        If you pass `generate_transcript=True` for an image with no text,
169        you don't get any alt text.
170        """
171        entity = create_image_entity(
172            fixtures_dir / "blue.png", generate_transcript=True
173        )
174        assert "alt_text" not in entity
175
176    def test_create_thumbnail_by_width(
177        self, fixtures_dir: Path, tmp_path: Path
178    ) -> None:
179        """
180        Create a thumbnail by width.
181        """
182        entity = create_image_entity(
183            fixtures_dir / "blue.png",
184            thumbnail_config={"out_dir": tmp_path / "thumbnails", "width": 10},
185        )
186
187        assert Path(entity["thumbnail_path"]).exists()
188
189        with Image.open(entity["thumbnail_path"]) as im:
190            assert im.width == 10
191
192    def test_create_thumbnail_by_height(
193        self, fixtures_dir: Path, tmp_path: Path
194    ) -> None:
195        """
196        Create a thumbnail by height.
197        """
198        entity = create_image_entity(
199            fixtures_dir / "blue.png",
200            thumbnail_config={"out_dir": tmp_path / "thumbnails", "height": 5},
201        )
202
203        assert Path(entity["thumbnail_path"]).exists()
204
205        with Image.open(entity["thumbnail_path"]) as im:
206            assert im.height == 5
207
208    @pytest.mark.parametrize(
209        "background, tint_colour",
210        [
211            ("white", "#005493"),
212            ("black", "#b3fdff"),
213            ("#111111", "#b3fdff"),
214        ],
215    )
216    def test_tint_colour_is_based_on_background(
217        self, fixtures_dir: Path, background: str, tint_colour: str
218    ) -> None:
219        """
220        The tint colour is based to suit the background.
221        """
222        # This is a checkerboard pattern made of 2 different shades of
223        # turquoise, a light and a dark.
224        entity = create_image_entity(
225            fixtures_dir / "checkerboard.png", background=background
226        )
227        assert entity["tint_colour"] == tint_colour
228
229
230class TestCreateVideoEntity:
231    """
232    Tests for `create_video_entity()`.
233    """
234
235    def test_basic_video(self, fixtures_dir: Path) -> None:
236        """
237        Get a video entity for a basic video.
238        """
239        # This video was downloaded from
240        # https://test-videos.co.uk/sintel/mp4-h264
241        entity = create_video_entity(
242            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
243            poster_path=fixtures_dir / "Sintel_360_10s_1MB_H264.png",
244        )
245        assert entity == {
246            "type": "video",
247            "path": "tests/fixtures/media/Sintel_360_10s_1MB_H264.mp4",
248            "width": 640,
249            "height": 360,
250            "duration": "0:00:10.000000",
251            "poster": {
252                "type": "image",
253                "path": "tests/fixtures/media/Sintel_360_10s_1MB_H264.png",
254                "tint_colour": "#020202",
255                "width": 640,
256                "height": 360,
257            },
258        }
259
260    def test_other_attrs_are_forwarded(self, fixtures_dir: Path) -> None:
261        """
262        The `subtitles_path`, `source_url` and `autoplay` values are
263        forwarded to the final entity.
264        """
265        entity = create_video_entity(
266            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
267            poster_path=fixtures_dir / "Sintel_360_10s_1MB_H264.png",
268            subtitles_path=fixtures_dir / "Sintel_360_10s_1MB_H264.en.vtt",
269            source_url="https://test-videos.co.uk/sintel/mp4-h264",
270            autoplay=True,
271        )
272
273        assert entity["subtitles"] == [
274            {
275                "path": "tests/fixtures/media/Sintel_360_10s_1MB_H264.en.vtt",
276                "label": "English",
277            }
278        ]
279        assert entity["source_url"] == "https://test-videos.co.uk/sintel/mp4-h264"
280        assert entity["autoplay"]
281
282    def test_gets_display_dimensions(self, fixtures_dir: Path) -> None:
283        """
284        The width/height dimensions are based on the display aspect ratio,
285        not the storage aspect ratio.
286
287        See https://alexwlchan.net/2025/square-pixels/
288        """
289        # This is a short clip of https://www.youtube.com/watch?v=HHhyznZ2u4E
290        entity = create_video_entity(
291            fixtures_dir / "Mars 2020 EDL Remastered [HHhyznZ2u4E].mp4",
292            poster_path=fixtures_dir / "Mars 2020 EDL Remastered [HHhyznZ2u4E].jpg",
293        )
294
295        assert entity["width"] == 1350
296        assert entity["height"] == 1080
297
298    def test_video_without_sample_aspect_ratio(self, fixtures_dir: Path) -> None:
299        """
300        Get the width/height dimensions of a video that doesn't have
301        `sample_aspect_ratio` in its metadata.
302        """
303        # This is a short clip from Wings (1927).
304        entity = create_video_entity(
305            fixtures_dir / "wings_tracking_shot.mp4",
306            poster_path=fixtures_dir / "wings_tracking_shot.jpg",
307        )
308
309        assert entity["width"] == 960
310        assert entity["height"] == 720
311
312    @pytest.mark.parametrize(
313        "background, tint_colour",
314        [
315            ("white", "#005493"),
316            ("black", "#b3fdff"),
317            ("#111111", "#b3fdff"),
318        ],
319    )
320    def test_tint_colour_is_based_on_background(
321        self, fixtures_dir: Path, background: str, tint_colour: str
322    ) -> None:
323        """
324        The tint colour is based to suit the background.
325        """
326        # The poster image is a checkerboard pattern made of 2 different
327        # shades of turquoise, a light and a dark.
328        entity = create_video_entity(
329            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
330            poster_path=fixtures_dir / "checkerboard.png",
331            background=background,
332        )
333        assert entity["poster"]["tint_colour"] == tint_colour
334
335    def test_video_with_thumbnail(self, fixtures_dir: Path, tmp_path: Path) -> None:
336        """
337        Create a low-resolution thumbnail of the poster image.
338        """
339        entity = create_video_entity(
340            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
341            poster_path=fixtures_dir / "Sintel_360_10s_1MB_H264.png",
342            thumbnail_config={"out_dir": tmp_path / "thumbnails", "width": 300},
343        )
344
345        assert entity["poster"]["thumbnail_path"] == str(
346            tmp_path / "thumbnails/Sintel_360_10s_1MB_H264.png"
347        )
348        assert Path(entity["poster"]["thumbnail_path"]).exists()
349
350
351class TestGetMediaPaths:
352    """
353    Tests for `get_media_paths`.
354    """
355
356    def test_basic_image(self, fixtures_dir: Path) -> None:
357        """
358        An image with no thumbnail only has one path: the image.
359        """
360        entity = create_image_entity(fixtures_dir / "blue.png")
361        assert get_media_paths(entity) == {fixtures_dir / "blue.png"}
362
363    def test_image_with_thumbnail(self, fixtures_dir: Path, tmp_path: Path) -> None:
364        """
365        An image with a thumbnail has two paths: the video and the
366        thumbnail.
367        """
368        entity = create_image_entity(
369            fixtures_dir / "blue.png",
370            thumbnail_config={"out_dir": tmp_path / "thumbnails", "width": 300},
371        )
372        assert get_media_paths(entity) == {
373            fixtures_dir / "blue.png",
374            tmp_path / "thumbnails/blue.png",
375        }
376
377    def test_video(self, fixtures_dir: Path) -> None:
378        """
379        A video has two paths: the video and the poster image.
380        """
381        entity = create_video_entity(
382            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
383            poster_path=fixtures_dir / "Sintel_360_10s_1MB_H264.png",
384        )
385        assert get_media_paths(entity) == {
386            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
387            fixtures_dir / "Sintel_360_10s_1MB_H264.png",
388        }
389
390    def test_video_with_subtitles(self, fixtures_dir: Path) -> None:
391        """
392        A video with subtitles has three paths: the video, the subtitles,
393        and the poster image.
394        """
395        entity = create_video_entity(
396            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
397            poster_path=fixtures_dir / "Sintel_360_10s_1MB_H264.png",
398            subtitles_path=fixtures_dir / "Sintel_360_10s_1MB_H264.en.vtt",
399        )
400        assert get_media_paths(entity) == {
401            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
402            fixtures_dir / "Sintel_360_10s_1MB_H264.png",
403            fixtures_dir / "Sintel_360_10s_1MB_H264.en.vtt",
404        }
405
406    def test_video_with_thumbnail(self, fixtures_dir: Path, tmp_path: Path) -> None:
407        """
408        A video with a poster thumbnail has three paths: the video,
409        the poster image, and the poster thumbnail.
410        """
411        entity = create_video_entity(
412            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
413            poster_path=fixtures_dir / "Sintel_360_10s_1MB_H264.png",
414            thumbnail_config={"out_dir": tmp_path / "thumbnails", "width": 300},
415        )
416        assert get_media_paths(entity) == {
417            fixtures_dir / "Sintel_360_10s_1MB_H264.mp4",
418            fixtures_dir / "Sintel_360_10s_1MB_H264.png",
419            tmp_path / "thumbnails/Sintel_360_10s_1MB_H264.png",
420        }
421
422    @pytest.mark.parametrize("bad_entity", [{}, {"type": "shape"}])
423    def test_unrecognised_entity_is_error(self, bad_entity: Any) -> None:
424        """
425        Getting media paths for an unrecognised entity type is a TypeError.
426        """
427        with pytest.raises(TypeError):
428            get_media_paths(bad_entity)
429
430
431class TestConvertSrtToVtt:
432    """
433    Tests for `convert_srt_to_vtt`.
434    """
435
436    def test_converting_srt_file(self, tmp_path: Path) -> None:
437        """
438        Check an SRT file is converted to VTT correctly.
439        """
440        srt_path = tmp_path / "example.en.srt"
441        srt_path.write_text(
442            "1\n"
443            "00:00:01,001 --> 00:00:10,010\n"
444            "Somebody said the first thing\n"
445            "\n"
446            "2\n"
447            "00:02:00,002 --> 00:20:00,020\n"
448            "Somebody else said the second thing\n"
449            "\n"
450            "3\n"
451            "03:00:00,003 --> 30:00:00,300\n"
452            "Yet another person said the third thing!"
453        )
454
455        vtt_path = convert_srt_to_vtt(srt_path)
456        assert vtt_path == (tmp_path / "example.en.vtt")
457        assert vtt_path.exists()
458        assert not srt_path.exists()
459
460        assert vtt_path.read_text() == (
461            "WEBVTT\n"
462            "\n"
463            "00:00:01.001 --> 00:00:10.010\n"
464            "Somebody said the first thing\n"
465            "\n"
466            "00:02:00.002 --> 00:20:00.020\n"
467            "Somebody else said the second thing\n"
468            "\n"
469            "03:00:00.003 --> 30:00:00.300\n"
470            "Yet another person said the third thing!"
471        )
472
473    def test_converting_srt_with_bom(self, tmp_path: Path) -> None:
474        """
475        If an SRT file starts with \ufeff (a byte order mark), it's
476        correctly removed from the caption number.
477        """
478        srt_path = tmp_path / "captions_with_bom.en.srt"
479        srt_path.write_text(
480            "\ufeff1\n"
481            "00:00:01,001 --> 00:00:10,010\n"
482            "Somebody said the first thing\n"
483            "\n"
484            "2\n"
485            "00:02:00,002 --> 00:20:00,020\n"
486            "Somebody else said the second thing\n"
487        )
488
489        vtt_path = convert_srt_to_vtt(srt_path)
490        assert vtt_path == (tmp_path / "captions_with_bom.en.vtt")
491        assert vtt_path.exists()
492        assert not srt_path.exists()
493
494        assert vtt_path.read_text() == (
495            "WEBVTT\n"
496            "\n"
497            "00:00:01.001 --> 00:00:10.010\n"
498            "Somebody said the first thing\n"
499            "\n"
500            "00:02:00.002 --> 00:20:00.020\n"
501            "Somebody else said the second thing\n"
502        )
503
504    def test_converting_srt_with_bom_on_later_number(self, tmp_path: Path) -> None:
505        """
506        If an SRT file starts with a byte order mark but the caption
507        number doesn't start at 1, the BOM is removed.
508        """
509        srt_path = tmp_path / "captions_with_bom.en.srt"
510        srt_path.write_text(
511            "\ufeff2\n"
512            "00:00:02,002 --> 00:00:20,020\n"
513            "Who said the first thing?\n"
514            "\n"
515            "3\n"
516            "00:03:00,003 --> 00:30:00,030\n"
517            "A ghost said the first thing\n"
518        )
519
520        vtt_path = convert_srt_to_vtt(srt_path)
521        assert vtt_path == (tmp_path / "captions_with_bom.en.vtt")
522        assert vtt_path.exists()
523        assert not srt_path.exists()
524
525        assert vtt_path.read_text() == (
526            "WEBVTT\n"
527            "\n"
528            "00:00:02.002 --> 00:00:20.020\n"
529            "Who said the first thing?\n"
530            "\n"
531            "00:03:00.003 --> 00:30:00.030\n"
532            "A ghost said the first thing\n"
533        )
534
535    def test_convert_empty_file_is_error(self, tmp_path: Path) -> None:
536        """
537        Converting an empty file is an error.
538        """
539        srt_path = tmp_path / "captions_with_bom.en.srt"
540        srt_path.write_text("")
541
542        with pytest.raises(ValueError, match="cannot convert empty SRT file"):
543            convert_srt_to_vtt(srt_path)
544
545    def test_converting_non_srt_file_is_error(self) -> None:
546        """
547        Converting a non-SRT file returns an error.
548        """
549        with pytest.raises(ValueError, match="can only take .srt files"):
550            convert_srt_to_vtt(srt_path=Path("example.txt"))