linkding/bookmarks/tests/test_preview_image_loader.py

import io
import os
import tempfile
from pathlib import Path
from unittest import mock

from django.conf import settings
from django.test import TestCase

from bookmarks.services import preview_image_loader

mock_image_data = b"mock_image"


class MockStreamingResponse:
    def __init__(
        self,
        url,
        data=mock_image_data,
        content_type="image/png",
        content_length=None,
        status_code=200,
    ):
        self.url = url
        self.chunks = [data]
        self.status_code = status_code
        if not content_length:
            content_length = len(data)
        self.headers = {"Content-Type": content_type, "Content-Length": content_length}

    def iter_content(self, **kwargs):
        return self.chunks

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass


class PreviewImageLoaderTestCase(TestCase):
    def setUp(self) -> None:
        self.temp_folder = tempfile.TemporaryDirectory()
        self.settings_override = self.settings(LD_PREVIEW_FOLDER=self.temp_folder.name)
        self.settings_override.enable()
        self.mock_load_website_metadata_patcher = mock.patch(
            "bookmarks.services.website_loader.load_website_metadata"
        )
        self.mock_load_website_metadata = (
            self.mock_load_website_metadata_patcher.start()
        )
        self.mock_load_website_metadata.return_value = mock.Mock(
            preview_image="https://example.com/image.png"
        )

    def tearDown(self) -> None:
        self.temp_folder.cleanup()
        self.settings_override.disable()
        self.mock_load_website_metadata_patcher.stop()

    def create_mock_response(
        self,
        url="https://example.com/image.png",
        icon_data=mock_image_data,
        content_type="image/png",
        content_length=len(mock_image_data),
        status_code=200,
    ):
        mock_response = mock.Mock()
        mock_response.raw = io.BytesIO(icon_data)
        return MockStreamingResponse(
            url, icon_data, content_type, content_length, status_code
        )

    def get_image_path(self, filename):
        return Path(os.path.join(settings.LD_PREVIEW_FOLDER, filename))

    def assertImageExists(self, filename, data):
        self.assertTrue(self.get_image_path(filename).exists())
        self.assertEqual(self.get_image_path(filename).read_bytes(), data)

    def assertNoImageExists(self):
        self.assertFalse(os.listdir(settings.LD_PREVIEW_FOLDER))

    def test_load_preview_image(self):
        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response()

            file = preview_image_loader.load_preview_image("https://example.com")

            self.assertIsNotNone(file)
            self.assertImageExists(file, mock_image_data)

    def test_load_preview_image_returns_none_if_no_preview_image_detected(self):
        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response()
            self.mock_load_website_metadata.return_value = mock.Mock(preview_image=None)

            file = preview_image_loader.load_preview_image("https://example.com")

            self.assertIsNone(file)
            self.assertNoImageExists()

    def test_load_preview_image_returns_none_for_invalid_status_code(self):
        invalid_status_codes = [199, 300, 400, 500]

        for status_code in invalid_status_codes:
            with mock.patch("requests.get") as mock_get:
                mock_get.return_value = self.create_mock_response(
                    status_code=status_code
                )

                file = preview_image_loader.load_preview_image("https://example.com")

                self.assertIsNone(file)
                self.assertNoImageExists()

    def test_load_preview_image_returns_none_if_content_length_exceeds_limit(self):
        # exceeds max size
        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response(
                content_length=settings.LD_PREVIEW_MAX_SIZE + 1
            )

            file = preview_image_loader.load_preview_image("https://example.com")

            self.assertIsNone(file)
            self.assertNoImageExists()

        # equals max size
        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response(
                content_length=settings.LD_PREVIEW_MAX_SIZE
            )

            file = preview_image_loader.load_preview_image("https://example.com")

            self.assertIsNotNone(file)
            self.assertImageExists(file, mock_image_data)

    def test_load_preview_image_returns_none_for_invalid_content_type(self):
        invalid_content_types = ["text/html", "application/json"]

        for content_type in invalid_content_types:
            with mock.patch("requests.get") as mock_get:
                mock_get.return_value = self.create_mock_response(
                    content_type=content_type
                )

                file = preview_image_loader.load_preview_image("https://example.com")

                self.assertIsNone(file)
                self.assertNoImageExists()

        valid_content_types = ["image/png", "image/jpeg", "image/gif"]

        for content_type in valid_content_types:
            with mock.patch("requests.get") as mock_get:
                mock_get.return_value = self.create_mock_response(
                    content_type=content_type
                )

                file = preview_image_loader.load_preview_image("https://example.com")

                self.assertIsNotNone(file)
                self.assertImageExists(file, mock_image_data)

    def test_load_preview_image_returns_none_if_download_exceeds_content_length(self):
        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response(content_length=1)

            file = preview_image_loader.load_preview_image("https://example.com")

            self.assertIsNone(file)
            self.assertNoImageExists()

    def test_load_preview_image_creates_folder_if_not_exists(self):
        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response()

            folder = Path(settings.LD_PREVIEW_FOLDER)
            folder.rmdir()

            self.assertFalse(folder.exists())

            preview_image_loader.load_preview_image("https://example.com")

            self.assertTrue(folder.exists())

    def test_guess_file_extension(self):
        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response(content_type="image/png")

            file = preview_image_loader.load_preview_image("https://example.com")

            self.assertImageExists(file, mock_image_data)
            self.assertEqual("png", file.split(".")[-1])

        with mock.patch("requests.get") as mock_get:
            mock_get.return_value = self.create_mock_response(content_type="image/jpeg")

            file = preview_image_loader.load_preview_image("https://example.com")

            self.assertImageExists(file, mock_image_data)
            self.assertEqual("jpg", file.split(".")[-1])