Merge branch 'master' into hunternif

4 files changed, 186 insertions, 20 deletions

diff --git a/server/szurubooru/tests/func/test_image_hash.py b/server/szurubooru/tests/func/test_image_hash.py
index e7028b6..5a5dc71 100644
--- a/server/szurubooru/tests/func/test_image_hash.py
+++ b/server/szurubooru/tests/func/test_image_hash.py
@@ -27,3 +27,53 @@ def test_signature_functions(read_asset, config_injector):
     words2 = image_hash.generate_words(sig2)
     words_match = sum(word1 == word2 for word1, word2 in zip(words1, words2))
     assert words_match == 18
+
+
+def test_signature_heif(read_asset, config_injector):
+    sig1 = image_hash.generate_signature(read_asset("heif.heif"))
+    sig2 = image_hash.generate_signature(read_asset("heif-similar.heif"))
+
+    sig1_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig1)
+    )
+    sig2_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig2)
+    )
+    assert array_equal(sig1, sig1_repacked)
+    assert array_equal(sig2, sig2_repacked)
+
+    dist1 = image_hash.normalized_distance([sig1], sig2)
+    assert abs(dist1[0] - 0.136777724290135) < 1e-8
+
+    dist2 = image_hash.normalized_distance([sig2], sig2)
+    assert abs(dist2[0]) < 1e-8
+
+    words1 = image_hash.generate_words(sig1)
+    words2 = image_hash.generate_words(sig2)
+    words_match = sum(word1 == word2 for word1, word2 in zip(words1, words2))
+    assert words_match == 43
+
+
+def test_signature_avif(read_asset, config_injector):
+    sig1 = image_hash.generate_signature(read_asset("avif.avif"))
+    sig2 = image_hash.generate_signature(read_asset("avif-similar.avif"))
+
+    sig1_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig1)
+    )
+    sig2_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig2)
+    )
+    assert array_equal(sig1, sig1_repacked)
+    assert array_equal(sig2, sig2_repacked)
+
+    dist1 = image_hash.normalized_distance([sig1], sig2)
+    assert abs(dist1[0] - 0.22628712858355998) < 1e-8
+
+    dist2 = image_hash.normalized_distance([sig2], sig2)
+    assert abs(dist2[0]) < 1e-8
+
+    words1 = image_hash.generate_words(sig1)
+    words2 = image_hash.generate_words(sig2)
+    words_match = sum(word1 == word2 for word1, word2 in zip(words1, words2))
+    assert words_match == 12
diff --git a/server/szurubooru/tests/func/test_mime.py b/server/szurubooru/tests/func/test_mime.py
index 0d8f645..b33746b 100644
--- a/server/szurubooru/tests/func/test_mime.py
+++ b/server/szurubooru/tests/func/test_mime.py
@@ -13,6 +13,12 @@ from szurubooru.func import mime
         ("jpeg.jpg", "image/jpeg"),
         ("gif.gif", "image/gif"),
         ("webp.webp", "image/webp"),
+        ("bmp.bmp", "image/bmp"),
+        ("avif.avif", "image/avif"),
+        ("avif-avis.avif", "image/avif"),
+        ("heif.heif", "image/heif"),
+        ("heic.heic", "image/heic"),
+        ("heic-heix.heic", "image/heic"),
         ("text.txt", "application/octet-stream"),
     ],
 )
@@ -34,6 +40,10 @@ def test_get_mime_type_for_empty_file():
         ("image/jpeg", "jpg"),
         ("image/gif", "gif"),
         ("image/webp", "webp"),
+        ("image/bmp", "bmp"),
+        ("image/avif", "avif"),
+        ("image/heif", "heif"),
+        ("image/heic", "heic"),
         ("application/octet-stream", "dat"),
     ],
 )
@@ -75,9 +85,17 @@ def test_is_video(input_mime_type, expected_state):
         ("image/gif", True),
         ("image/png", True),
         ("image/jpeg", True),
+        ("image/bmp", True),
+        ("image/avif", True),
+        ("image/heic", True),
+        ("image/heif", True),
         ("IMAGE/GIF", True),
         ("IMAGE/PNG", True),
         ("IMAGE/JPEG", True),
+        ("IMAGE/BMP", True),
+        ("IMAGE/AVIF", True),
+        ("IMAGE/HEIC", True),
+        ("IMAGE/HEIF", True),
         ("image/anything_else", False),
         ("not an image", False),
     ],
@@ -95,3 +113,26 @@ def test_is_image(input_mime_type, expected_state):
 )
 def test_is_animated_gif(read_asset, input_path, expected_state):
     assert mime.is_animated_gif(read_asset(input_path)) == expected_state
+
+
+@pytest.mark.parametrize(
+    "input_mime_type,expected_state",
+    [
+        ("image/gif", False),
+        ("image/png", False),
+        ("image/jpeg", False),
+        ("image/avif", True),
+        ("image/heic", True),
+        ("image/heif", True),
+        ("IMAGE/GIF", False),
+        ("IMAGE/PNG", False),
+        ("IMAGE/JPEG", False),
+        ("IMAGE/AVIF", True),
+        ("IMAGE/HEIC", True),
+        ("IMAGE/HEIF", True),
+        ("image/anything_else", False),
+        ("not an image", False),
+    ],
+)
+def test_is_heif(input_mime_type, expected_state):
+    assert mime.is_heif(input_mime_type) == expected_state
diff --git a/server/szurubooru/tests/func/test_net.py b/server/szurubooru/tests/func/test_net.py
index 65e9048..c5b4c73 100644
--- a/server/szurubooru/tests/func/test_net.py
+++ b/server/szurubooru/tests/func/test_net.py
@@ -1,6 +1,3 @@
-from datetime import datetime
-from unittest.mock import patch
-
 import pytest
 
 from szurubooru import errors
@@ -69,40 +66,38 @@ def test_download():
     "url",
     [
         "https://samples.ffmpeg.org/MPEG-4/video.mp4",
+        "https://www.youtube.com/watch?v=dQw4w9WgXcQ",
     ],
 )
 def test_too_large_download(url):
-    pytest.xfail("Download limit not implemented yet")
-    with pytest.raises(errors.ProcessingError):
-        net.download(url)
+    with pytest.raises(net.DownloadTooLargeError):
+        net.download(url, use_video_downloader=True)
 
 
 @pytest.mark.parametrize(
     "url,expected_sha1",
     [
         (
-            "https://www.youtube.com/watch?v=C0DPdy98e4c",
-            "365af1c8f59c6865e1a84c6e13e3e25ff89e0ba1",
+            "https://gfycat.com/immaterialchillyiberianmole",
+            "0125976d2439e651b6863438db30de58f79f7754",
+        ),
+        (
+            "https://upload.wikimedia.org/wikipedia/commons/a/ad/Utah_teapot.png",  # noqa: E501
+            "cfadcbdeda1204dc1363ee5c1969191f26be2e41",
         ),
         (
-            "https://gfycat.com/immaterialchillyiberianmole",
-            "953000e81d7bd1da95ce264f872e7b6c4a6484be",
+            "https://i.imgur.com/GPgh0AN.jpg",
+            "26861a4663fedae48e5beed3eec5156ded20640f",
         ),
     ],
 )
-def test_video_download(url, expected_sha1):
+def test_content_download(url, expected_sha1):
     actual_content = net.download(url, use_video_downloader=True)
     assert get_sha1(actual_content) == expected_sha1
 
 
-@pytest.mark.parametrize(
-    "url",
-    [
-        "https://samples.ffmpeg.org/flac/short.flac",  # not a video
-        "https://www.youtube.com/watch?v=dQw4w9WgXcQ",  # video too large
-    ],
-)
-def test_failed_video_download(url):
+def test_bad_content_downlaod():
+    url = "http://info.cern.ch/hypertext/WWW/TheProject.html"
     with pytest.raises(errors.ThirdPartyError):
         net.download(url, use_video_downloader=True)
 
diff --git a/server/szurubooru/tests/func/test_posts.py b/server/szurubooru/tests/func/test_posts.py
index af6121a..e1be764 100644
--- a/server/szurubooru/tests/func/test_posts.py
+++ b/server/szurubooru/tests/func/test_posts.py
@@ -147,6 +147,7 @@ def test_serialize_post(
         post.source = "4gag"
         post.type = model.Post.TYPE_IMAGE
         post.checksum = "deadbeef"
+        post.checksum_md5 = "deadbeef"
         post.mime_type = "image/jpeg"
         post.file_size = 100
         post.user = user_factory(name="post author")
@@ -231,6 +232,7 @@ def test_serialize_post(
             "source": "4gag",
             "type": "image",
             "checksum": "deadbeef",
+            "checksumMD5": "deadbeef",
             "fileSize": 100,
             "canvasWidth": 200,
             "canvasHeight": 300,
@@ -424,6 +426,48 @@ def test_update_post_source_with_too_long_string():
         ),
         (
             False,
+            "bmp.bmp",
+            "image/bmp",
+            model.Post.TYPE_IMAGE,
+            "1_244c8840887984c4.bmp",
+        ),
+        (
+            False,
+            "avif.avif",
+            "image/avif",
+            model.Post.TYPE_IMAGE,
+            "1_244c8840887984c4.avif",
+        ),
+        (
+            False,
+            "avif-avis.avif",
+            "image/avif",
+            model.Post.TYPE_IMAGE,
+            "1_244c8840887984c4.avif",
+        ),
+        (
+            False,
+            "heic.heic",
+            "image/heic",
+            model.Post.TYPE_IMAGE,
+            "1_244c8840887984c4.heic",
+        ),
+        (
+            False,
+            "heic-heix.heic",
+            "image/heic",
+            model.Post.TYPE_IMAGE,
+            "1_244c8840887984c4.heic",
+        ),
+        (
+            False,
+            "heif.heif",
+            "image/heif",
+            model.Post.TYPE_IMAGE,
+            "1_244c8840887984c4.heif",
+        ),
+        (
+            False,
             "gif-animated.gif",
             "image/gif",
             model.Post.TYPE_ANIMATION,
@@ -463,8 +507,11 @@ def test_update_post_content_for_new_post(
     expected_type,
     output_file_name,
 ):
-    with patch("szurubooru.func.util.get_sha1"):
+    with patch("szurubooru.func.util.get_sha1"), patch(
+        "szurubooru.func.util.get_md5"
+    ):
         util.get_sha1.return_value = "crc"
+        util.get_md5.return_value = "md5"
         config_injector(
             {
                 "data_dir": str(tmpdir.mkdir("data")),
@@ -490,6 +537,7 @@ def test_update_post_content_for_new_post(
         assert post.mime_type == expected_mime_type
         assert post.type == expected_type
         assert post.checksum == "crc"
+        assert post.checksum_md5 == "md5"
         assert os.path.exists(output_file_path)
         if post.type in (model.Post.TYPE_IMAGE, model.Post.TYPE_ANIMATION):
             assert db.session.query(model.PostSignature).count() == 1
@@ -725,6 +773,38 @@ def test_update_post_content_leaving_custom_thumbnail(
     assert os.path.exists(generated_path)
 
 
+@pytest.mark.parametrize("filename", ("avif.avif", "heic.heic", "heif.heif"))
+def test_update_post_content_convert_heif_to_png_when_processing(
+    tmpdir, config_injector, read_asset, post_factory, filename
+):
+    config_injector(
+        {
+            "data_dir": str(tmpdir.mkdir("data")),
+            "thumbnails": {
+                "post_width": 300,
+                "post_height": 300,
+            },
+            "secret": "test",
+            "allow_broken_uploads": False,
+        }
+    )
+    post = post_factory(id=1)
+    db.session.add(post)
+    posts.update_post_content(post, read_asset(filename))
+    posts.update_post_thumbnail(post, read_asset(filename))
+    db.session.flush()
+    generated_path = (
+        "{}/data/generated-thumbnails/".format(tmpdir)
+        + "1_244c8840887984c4.jpg"
+    )
+    source_path = (
+        "{}/data/posts/custom-thumbnails/".format(tmpdir)
+        + "1_244c8840887984c4.dat"
+    )
+    assert os.path.exists(source_path)
+    assert os.path.exists(generated_path)
+
+
 def test_update_post_tags(tag_factory):
     post = model.Post()
     with patch("szurubooru.func.tags.get_or_create_tags_by_names"):