linkding/bookmarks/tests/test_bookmarks_tasks.py

import os.path
from unittest import mock

import waybackpy
from django.conf import settings
from django.contrib.auth.models import User
from django.test import TestCase, override_settings
from huey.contrib.djhuey import HUEY as huey
from waybackpy.exceptions import WaybackError

from bookmarks.models import BookmarkAsset, UserProfile
from bookmarks.services import tasks, singlefile
from bookmarks.tests.helpers import BookmarkFactoryMixin


def create_wayback_machine_save_api_mock(
    archive_url: str = "https://example.com/created_snapshot",
    fail_on_save: bool = False,
):
    mock_api = mock.Mock(archive_url=archive_url)

    if fail_on_save:
        mock_api.save.side_effect = WaybackError

    return mock_api


class BookmarkTasksTestCase(TestCase, BookmarkFactoryMixin):

    def setUp(self):
        huey.immediate = True
        huey.results = True
        huey.store_none = True

        self.mock_save_api = mock.Mock(
            archive_url="https://example.com/created_snapshot"
        )
        self.mock_save_api_patcher = mock.patch.object(
            waybackpy, "WaybackMachineSaveAPI", return_value=self.mock_save_api
        )
        self.mock_save_api_patcher.start()

        self.mock_load_favicon_patcher = mock.patch(
            "bookmarks.services.favicon_loader.load_favicon"
        )
        self.mock_load_favicon = self.mock_load_favicon_patcher.start()
        self.mock_load_favicon.return_value = "https_example_com.png"

        self.mock_singlefile_create_snapshot_patcher = mock.patch(
            "bookmarks.services.singlefile.create_snapshot",
        )
        self.mock_singlefile_create_snapshot = (
            self.mock_singlefile_create_snapshot_patcher.start()
        )

        self.mock_load_preview_image_patcher = mock.patch(
            "bookmarks.services.preview_image_loader.load_preview_image"
        )
        self.mock_load_preview_image = self.mock_load_preview_image_patcher.start()
        self.mock_load_preview_image.return_value = "preview_image.png"

        user = self.get_or_create_test_user()
        user.profile.web_archive_integration = (
            UserProfile.WEB_ARCHIVE_INTEGRATION_ENABLED
        )
        user.profile.enable_favicons = True
        user.profile.enable_preview_images = True
        user.profile.save()

    def tearDown(self):
        self.mock_save_api_patcher.stop()
        self.mock_load_favicon_patcher.stop()
        self.mock_singlefile_create_snapshot_patcher.stop()
        self.mock_load_preview_image_patcher.stop()
        huey.storage.flush_results()
        huey.immediate = False

    def executed_count(self):
        return len(huey.all_results())

    def test_create_web_archive_snapshot_should_update_snapshot_url(self):
        bookmark = self.setup_bookmark()

        tasks.create_web_archive_snapshot(
            self.get_or_create_test_user(), bookmark, False
        )
        bookmark.refresh_from_db()

        self.mock_save_api.save.assert_called_once()
        self.assertEqual(self.executed_count(), 1)
        self.assertEqual(
            bookmark.web_archive_snapshot_url,
            "https://example.com/created_snapshot",
        )

    def test_create_web_archive_snapshot_should_handle_missing_bookmark_id(self):
        tasks._create_web_archive_snapshot_task(123, False)

        self.assertEqual(self.executed_count(), 1)
        self.mock_save_api.save.assert_not_called()

    def test_create_web_archive_snapshot_should_skip_if_snapshot_exists(self):
        bookmark = self.setup_bookmark(web_archive_snapshot_url="https://example.com")

        self.mock_save_api.create_web_archive_snapshot(
            self.get_or_create_test_user(), bookmark, False
        )

        self.assertEqual(self.executed_count(), 0)
        self.mock_save_api.assert_not_called()

    def test_create_web_archive_snapshot_should_force_update_snapshot(self):
        bookmark = self.setup_bookmark(web_archive_snapshot_url="https://example.com")
        self.mock_save_api.archive_url = "https://other.com"

        tasks.create_web_archive_snapshot(
            self.get_or_create_test_user(), bookmark, True
        )
        bookmark.refresh_from_db()

        self.assertEqual(bookmark.web_archive_snapshot_url, "https://other.com")

    def test_create_web_archive_snapshot_should_not_save_stale_bookmark_data(self):
        bookmark = self.setup_bookmark()

        # update bookmark during API call to check that saving
        # the snapshot does not overwrite updated bookmark data
        def mock_save_impl():
            bookmark.title = "Updated title"
            bookmark.save()

        self.mock_save_api.save.side_effect = mock_save_impl

        tasks.create_web_archive_snapshot(
            self.get_or_create_test_user(), bookmark, False
        )
        bookmark.refresh_from_db()

        self.assertEqual(bookmark.title, "Updated title")
        self.assertEqual(
            "https://example.com/created_snapshot",
            bookmark.web_archive_snapshot_url,
        )

    @override_settings(LD_DISABLE_BACKGROUND_TASKS=True)
    def test_create_web_archive_snapshot_should_not_run_when_background_tasks_are_disabled(
        self,
    ):
        bookmark = self.setup_bookmark()

        tasks.create_web_archive_snapshot(
            self.get_or_create_test_user(), bookmark, False
        )
        self.assertEqual(self.executed_count(), 0)

    def test_create_web_archive_snapshot_should_not_run_when_web_archive_integration_is_disabled(
        self,
    ):
        self.user.profile.web_archive_integration = (
            UserProfile.WEB_ARCHIVE_INTEGRATION_DISABLED
        )
        self.user.profile.save()

        bookmark = self.setup_bookmark()
        tasks.create_web_archive_snapshot(
            self.get_or_create_test_user(), bookmark, False
        )

        self.assertEqual(self.executed_count(), 0)

    def test_load_favicon_should_create_favicon_file(self):
        bookmark = self.setup_bookmark()

        tasks.load_favicon(self.get_or_create_test_user(), bookmark)
        bookmark.refresh_from_db()

        self.assertEqual(self.executed_count(), 1)
        self.assertEqual(bookmark.favicon_file, "https_example_com.png")

    def test_load_favicon_should_update_favicon_file(self):
        bookmark = self.setup_bookmark(favicon_file="https_example_com.png")

        self.mock_load_favicon.return_value = "https_example_updated_com.png"

        tasks.load_favicon(self.get_or_create_test_user(), bookmark)

        bookmark.refresh_from_db()
        self.mock_load_favicon.assert_called_once()
        self.assertEqual(bookmark.favicon_file, "https_example_updated_com.png")

    def test_load_favicon_should_handle_missing_bookmark(self):
        tasks._load_favicon_task(123)

        self.mock_load_favicon.assert_not_called()

    def test_load_favicon_should_not_save_stale_bookmark_data(self):
        bookmark = self.setup_bookmark()

        # update bookmark during API call to check that saving
        # the favicon does not overwrite updated bookmark data
        def mock_load_favicon_impl(url):
            bookmark.title = "Updated title"
            bookmark.save()
            return "https_example_com.png"

        self.mock_load_favicon.side_effect = mock_load_favicon_impl

        tasks.load_favicon(self.get_or_create_test_user(), bookmark)
        bookmark.refresh_from_db()

        self.assertEqual(bookmark.title, "Updated title")
        self.assertEqual(bookmark.favicon_file, "https_example_com.png")

    @override_settings(LD_DISABLE_BACKGROUND_TASKS=True)
    def test_load_favicon_should_not_run_when_background_tasks_are_disabled(self):
        bookmark = self.setup_bookmark()
        tasks.load_favicon(self.get_or_create_test_user(), bookmark)

        self.assertEqual(self.executed_count(), 0)

    def test_load_favicon_should_not_run_when_favicon_feature_is_disabled(self):
        self.user.profile.enable_favicons = False
        self.user.profile.save()

        bookmark = self.setup_bookmark()
        tasks.load_favicon(self.get_or_create_test_user(), bookmark)

        self.assertEqual(self.executed_count(), 0)

    def test_schedule_bookmarks_without_favicons_should_load_favicon_for_all_bookmarks_without_favicon(
        self,
    ):
        user = self.get_or_create_test_user()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark(favicon_file="https_example_com.png")
        self.setup_bookmark(favicon_file="https_example_com.png")
        self.setup_bookmark(favicon_file="https_example_com.png")

        tasks.schedule_bookmarks_without_favicons(user)

        self.assertEqual(self.executed_count(), 4)
        self.assertEqual(self.mock_load_favicon.call_count, 3)

    def test_schedule_bookmarks_without_favicons_should_only_update_user_owned_bookmarks(
        self,
    ):
        user = self.get_or_create_test_user()
        other_user = User.objects.create_user(
            "otheruser", "otheruser@example.com", "password123"
        )
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)

        tasks.schedule_bookmarks_without_favicons(user)

        self.assertEqual(self.mock_load_favicon.call_count, 3)

    @override_settings(LD_DISABLE_BACKGROUND_TASKS=True)
    def test_schedule_bookmarks_without_favicons_should_not_run_when_background_tasks_are_disabled(
        self,
    ):
        self.setup_bookmark()
        tasks.schedule_bookmarks_without_favicons(self.get_or_create_test_user())

        self.assertEqual(self.executed_count(), 0)

    def test_schedule_bookmarks_without_favicons_should_not_run_when_favicon_feature_is_disabled(
        self,
    ):
        self.user.profile.enable_favicons = False
        self.user.profile.save()

        self.setup_bookmark()
        tasks.schedule_bookmarks_without_favicons(self.get_or_create_test_user())

        self.assertEqual(self.executed_count(), 0)

    def test_schedule_refresh_favicons_should_update_favicon_for_all_bookmarks(self):
        user = self.get_or_create_test_user()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark(favicon_file="https_example_com.png")
        self.setup_bookmark(favicon_file="https_example_com.png")
        self.setup_bookmark(favicon_file="https_example_com.png")

        tasks.schedule_refresh_favicons(user)

        self.assertEqual(self.executed_count(), 7)
        self.assertEqual(self.mock_load_favicon.call_count, 6)

    def test_schedule_refresh_favicons_should_only_update_user_owned_bookmarks(self):
        user = self.get_or_create_test_user()
        other_user = User.objects.create_user(
            "otheruser", "otheruser@example.com", "password123"
        )
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)

        tasks.schedule_refresh_favicons(user)

        self.assertEqual(self.mock_load_favicon.call_count, 3)

    @override_settings(LD_DISABLE_BACKGROUND_TASKS=True)
    def test_schedule_refresh_favicons_should_not_run_when_background_tasks_are_disabled(
        self,
    ):
        self.setup_bookmark()
        tasks.schedule_refresh_favicons(self.get_or_create_test_user())

        self.assertEqual(self.executed_count(), 0)

    @override_settings(LD_ENABLE_REFRESH_FAVICONS=False)
    def test_schedule_refresh_favicons_should_not_run_when_refresh_is_disabled(self):
        self.setup_bookmark()
        tasks.schedule_refresh_favicons(self.get_or_create_test_user())

        self.assertEqual(self.executed_count(), 0)

    def test_schedule_refresh_favicons_should_not_run_when_favicon_feature_is_disabled(
        self,
    ):
        self.user.profile.enable_favicons = False
        self.user.profile.save()

        self.setup_bookmark()
        tasks.schedule_refresh_favicons(self.get_or_create_test_user())

        self.assertEqual(self.executed_count(), 0)

    def test_load_preview_image_should_create_preview_image_file(self):
        bookmark = self.setup_bookmark()

        tasks.load_preview_image(self.get_or_create_test_user(), bookmark)
        bookmark.refresh_from_db()

        self.assertEqual(self.executed_count(), 1)
        self.assertEqual(bookmark.preview_image_file, "preview_image.png")

    def test_load_preview_image_should_update_preview_image_file(self):
        bookmark = self.setup_bookmark(
            preview_image_file="preview_image.png",
        )

        self.mock_load_preview_image.return_value = "preview_image_upd.png"

        tasks.load_preview_image(self.get_or_create_test_user(), bookmark)

        bookmark.refresh_from_db()
        self.mock_load_preview_image.assert_called_once()
        self.assertEqual(bookmark.preview_image_file, "preview_image_upd.png")

    def test_load_preview_image_should_set_blank_when_none_is_returned(self):
        bookmark = self.setup_bookmark(
            preview_image_file="preview_image.png",
        )

        self.mock_load_preview_image.return_value = None

        tasks.load_preview_image(self.get_or_create_test_user(), bookmark)

        bookmark.refresh_from_db()
        self.mock_load_preview_image.assert_called_once()
        self.assertEqual(bookmark.preview_image_file, "")

    def test_load_preview_image_should_handle_missing_bookmark(self):
        tasks._load_preview_image_task(123)

        self.mock_load_preview_image.assert_not_called()

    def test_load_preview_image_should_not_save_stale_bookmark_data(self):
        bookmark = self.setup_bookmark()

        # update bookmark during API call to check that saving
        # the image does not overwrite updated bookmark data
        def mock_load_preview_image_impl(url):
            bookmark.title = "Updated title"
            bookmark.save()
            return "test.png"

        self.mock_load_preview_image.side_effect = mock_load_preview_image_impl

        tasks.load_preview_image(self.get_or_create_test_user(), bookmark)
        bookmark.refresh_from_db()

        self.assertEqual(bookmark.title, "Updated title")
        self.assertEqual(bookmark.preview_image_file, "test.png")

    @override_settings(LD_DISABLE_BACKGROUND_TASKS=True)
    def test_load_preview_image_should_not_run_when_background_tasks_are_disabled(self):
        bookmark = self.setup_bookmark()
        tasks.load_preview_image(self.get_or_create_test_user(), bookmark)

        self.assertEqual(self.executed_count(), 0)

    def test_load_preview_image_should_not_run_when_preview_image_feature_is_disabled(
        self,
    ):
        self.user.profile.enable_preview_images = False
        self.user.profile.save()

        bookmark = self.setup_bookmark()
        tasks.load_preview_image(self.get_or_create_test_user(), bookmark)

        self.assertEqual(self.executed_count(), 0)

    def test_schedule_bookmarks_without_previews_should_load_preview_for_all_bookmarks_without_preview(
        self,
    ):
        user = self.get_or_create_test_user()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark(preview_image_file="test.png")
        self.setup_bookmark(preview_image_file="test.png")
        self.setup_bookmark(preview_image_file="test.png")

        tasks.schedule_bookmarks_without_previews(user)

        self.assertEqual(self.executed_count(), 4)
        self.assertEqual(self.mock_load_preview_image.call_count, 3)

    def test_schedule_bookmarks_without_previews_should_only_update_user_owned_bookmarks(
        self,
    ):
        user = self.get_or_create_test_user()
        other_user = User.objects.create_user(
            "otheruser", "otheruser@example.com", "password123"
        )
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)

        tasks.schedule_bookmarks_without_previews(user)

        self.assertEqual(self.mock_load_preview_image.call_count, 3)

    @override_settings(LD_DISABLE_BACKGROUND_TASKS=True)
    def test_schedule_bookmarks_without_previews_should_not_run_when_background_tasks_are_disabled(
        self,
    ):
        self.setup_bookmark()
        tasks.schedule_bookmarks_without_previews(self.get_or_create_test_user())

        self.assertEqual(self.executed_count(), 0)

    def test_schedule_bookmarks_without_previews_should_not_run_when_preview_feature_is_disabled(
        self,
    ):
        self.user.profile.enable_preview_images = False
        self.user.profile.save()

        self.setup_bookmark()
        tasks.schedule_bookmarks_without_previews(self.get_or_create_test_user())

        self.assertEqual(self.executed_count(), 0)

    @override_settings(LD_ENABLE_SNAPSHOTS=True)
    def test_create_html_snapshot_should_create_pending_asset(self):
        bookmark = self.setup_bookmark()

        # Mock the task function to avoid running it immediately
        with mock.patch("bookmarks.services.tasks._create_html_snapshot_task"):
            tasks.create_html_snapshot(bookmark)
            self.assertEqual(BookmarkAsset.objects.count(), 1)

            tasks.create_html_snapshot(bookmark)
            self.assertEqual(BookmarkAsset.objects.count(), 2)

            assets = BookmarkAsset.objects.filter(bookmark=bookmark)
            for asset in assets:
                self.assertEqual(asset.bookmark, bookmark)
                self.assertEqual(asset.asset_type, BookmarkAsset.TYPE_SNAPSHOT)
                self.assertEqual(asset.content_type, BookmarkAsset.CONTENT_TYPE_HTML)
                self.assertIn("HTML snapshot", asset.display_name)
                self.assertEqual(asset.status, BookmarkAsset.STATUS_PENDING)

    @override_settings(LD_ENABLE_SNAPSHOTS=True)
    def test_create_html_snapshot_should_update_file_info(self):
        bookmark = self.setup_bookmark(url="https://example.com")

        with mock.patch(
            "bookmarks.services.tasks._generate_snapshot_filename"
        ) as mock_generate:
            expected_filename = "snapshot_2021-01-02_034455_https___example.com.html.gz"
            mock_generate.return_value = expected_filename

            tasks.create_html_snapshot(bookmark)
            BookmarkAsset.objects.get(bookmark=bookmark)

            # Run periodic task to process the snapshot
            tasks._schedule_html_snapshots_task()

            self.mock_singlefile_create_snapshot.assert_called_once_with(
                "https://example.com",
                os.path.join(
                    settings.LD_ASSET_FOLDER,
                    expected_filename,
                ),
            )

            asset = BookmarkAsset.objects.get(bookmark=bookmark)
            self.assertEqual(asset.status, BookmarkAsset.STATUS_COMPLETE)
            self.assertEqual(asset.file, expected_filename)
            self.assertTrue(asset.gzip)

    @override_settings(LD_ENABLE_SNAPSHOTS=True)
    def test_create_html_snapshot_truncate_filename(self):
        # Create a bookmark with a very long URL
        long_url = "http://" + "a" * 300 + ".com"
        bookmark = self.setup_bookmark(url=long_url)

        tasks.create_html_snapshot(bookmark)
        BookmarkAsset.objects.get(bookmark=bookmark)

        # Run periodic task to process the snapshot
        tasks._schedule_html_snapshots_task()

        asset = BookmarkAsset.objects.get(bookmark=bookmark)
        self.assertEqual(len(asset.file), 192)

    @override_settings(LD_ENABLE_SNAPSHOTS=True)
    def test_create_html_snapshot_should_handle_error(self):
        bookmark = self.setup_bookmark(url="https://example.com")

        self.mock_singlefile_create_snapshot.side_effect = singlefile.SingleFileError(
            "Error"
        )
        tasks.create_html_snapshot(bookmark)

        # Run periodic task to process the snapshot
        tasks._schedule_html_snapshots_task()

        asset = BookmarkAsset.objects.get(bookmark=bookmark)
        self.assertEqual(asset.status, BookmarkAsset.STATUS_FAILURE)
        self.assertEqual(asset.file, "")
        self.assertFalse(asset.gzip)

    @override_settings(LD_ENABLE_SNAPSHOTS=True)
    def test_create_html_snapshot_should_handle_missing_bookmark(self):
        tasks._create_html_snapshot_task(123)

        self.mock_singlefile_create_snapshot.assert_not_called()

    @override_settings(LD_ENABLE_SNAPSHOTS=False)
    def test_create_html_snapshot_should_not_create_asset_when_single_file_is_disabled(
        self,
    ):
        bookmark = self.setup_bookmark()
        tasks.create_html_snapshot(bookmark)

        self.assertEqual(BookmarkAsset.objects.count(), 0)

    @override_settings(LD_ENABLE_SNAPSHOTS=True, LD_DISABLE_BACKGROUND_TASKS=True)
    def test_create_html_snapshot_should_not_create_asset_when_background_tasks_are_disabled(
        self,
    ):
        bookmark = self.setup_bookmark()
        tasks.create_html_snapshot(bookmark)

        self.assertEqual(BookmarkAsset.objects.count(), 0)

    @override_settings(LD_ENABLE_SNAPSHOTS=True)
    def test_create_missing_html_snapshots(self):
        bookmarks_with_snapshots = []
        bookmarks_without_snapshots = []

        # setup bookmarks with snapshots
        bookmark = self.setup_bookmark()
        self.setup_asset(
            bookmark=bookmark,
            asset_type=BookmarkAsset.TYPE_SNAPSHOT,
            status=BookmarkAsset.STATUS_COMPLETE,
        )
        bookmarks_with_snapshots.append(bookmark)

        bookmark = self.setup_bookmark()
        self.setup_asset(
            bookmark=bookmark,
            asset_type=BookmarkAsset.TYPE_SNAPSHOT,
            status=BookmarkAsset.STATUS_PENDING,
        )
        bookmarks_with_snapshots.append(bookmark)

        # setup bookmarks without snapshots
        bookmark = self.setup_bookmark()
        bookmarks_without_snapshots.append(bookmark)

        bookmark = self.setup_bookmark()
        self.setup_asset(
            bookmark=bookmark,
            asset_type=BookmarkAsset.TYPE_SNAPSHOT,
            status=BookmarkAsset.STATUS_FAILURE,
        )
        bookmarks_without_snapshots.append(bookmark)

        bookmark = self.setup_bookmark()
        self.setup_asset(
            bookmark=bookmark,
            asset_type="some_other_type",
            status=BookmarkAsset.STATUS_PENDING,
        )
        bookmarks_without_snapshots.append(bookmark)

        bookmark = self.setup_bookmark()
        self.setup_asset(
            bookmark=bookmark,
            asset_type="some_other_type",
            status=BookmarkAsset.STATUS_COMPLETE,
        )
        bookmarks_without_snapshots.append(bookmark)

        initial_assets = list(BookmarkAsset.objects.all())
        initial_assets_count = len(initial_assets)
        initial_asset_ids = [asset.id for asset in initial_assets]
        count = tasks.create_missing_html_snapshots(self.get_or_create_test_user())

        self.assertEqual(count, 4)
        self.assertEqual(BookmarkAsset.objects.count(), initial_assets_count + count)

        for bookmark in bookmarks_without_snapshots:
            new_assets = BookmarkAsset.objects.filter(bookmark=bookmark).exclude(
                id__in=initial_asset_ids
            )
            self.assertEqual(new_assets.count(), 1)

        for bookmark in bookmarks_with_snapshots:
            new_assets = BookmarkAsset.objects.filter(bookmark=bookmark).exclude(
                id__in=initial_asset_ids
            )
            self.assertEqual(new_assets.count(), 0)

    @override_settings(LD_ENABLE_SNAPSHOTS=True)
    def test_create_missing_html_snapshots_respects_current_user(self):
        self.setup_bookmark()
        self.setup_bookmark()
        self.setup_bookmark()

        other_user = self.setup_user()
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)
        self.setup_bookmark(user=other_user)

        count = tasks.create_missing_html_snapshots(self.get_or_create_test_user())

        self.assertEqual(count, 3)
        self.assertEqual(BookmarkAsset.objects.count(), count)