Improve tests

wagtail_localize/models.py

@@ -234,6 +234,14 @@ class CannotSaveDraftError(Exception):
     pass
 
 
+class NoViewRestrictionsError(Exception):
+    """
+    Raised when trying to sync view restrictions for non-Page objects
+    """
+
+    pass
+
+
 class MissingTranslationError(Exception):
     def __init__(self, segment, locale):
         self.segment = segment
@@ -888,7 +896,7 @@ def sync_view_restrictions(self, original, translation_page):
             translation_page (Page|Snippet): The translated instance.
         """
         if not isinstance(original, Page) or not isinstance(translation_page, Page):
-            return
+            raise NoViewRestrictionsError
 
         if original.view_restrictions.exists():
             original_restriction = original.view_restrictions.first()
@@ -943,7 +951,7 @@ def update_target_view_restrictions(self, locale):
 
         try:
             translation_page = self.get_translated_instance(locale)
-        except models.ObjectDoesNotExist:
+        except Page.DoesNotExist:
             return
 
         self.sync_view_restrictions(original, translation_page)

wagtail_localize/tests/test_translation_model.py

@@ -1,3 +1,5 @@
+from unittest import mock
+
 import polib
 
 from django.test import TestCase
@@ -6,6 +8,7 @@
 
 from wagtail_localize.models import (
     CannotSaveDraftError,
+    NoViewRestrictionsError,
     String,
     StringNotUsedInContext,
     StringTranslation,
@@ -748,3 +751,47 @@ def test_snippet(self):
         es_translation.refresh_from_db()
         self.assertFalse(fr_translation.enabled)
         self.assertTrue(es_translation.enabled)
+
+
+class TestTranslationSourceViewRestrictions(TestCase):
+    def setUp(self):
+        self.fr_locale = Locale.objects.create(language_code="fr")
+
+        self.page = create_test_page(
+            title="Test page",
+            slug="test-slug",
+            test_charfield="Test content",
+            test_textfield="More test content",
+        )
+        self.source = TranslationSource.objects.get()
+        self.translation = Translation.objects.create(
+            source=self.source,
+            target_locale=self.fr_locale,
+        )
+
+    @mock.patch("wagtail_localize.models.TranslationSource.sync_view_restrictions")
+    @mock.patch.object(
+        TranslationSource, "get_translated_instance", side_effect=Page.DoesNotExist()
+    )
+    def test_update_target_view_restrictions_with_missing_translation(
+        self, get_translated_instance, sync_view_restrictions
+    ):
+        self.source.update_target_view_restrictions(self.fr_locale)
+
+        self.assertEqual(get_translated_instance.call_count, 1)
+        self.assertEqual(sync_view_restrictions.call_count, 0)
+
+    @mock.patch("wagtail_localize.models.TranslationSource.get_translated_instance")
+    def test_update_target_restrictions_ignores_snippets(self, get_translated_instance):
+        snippet = TestSnippet.objects.create(field="Test content")
+        source, created = TranslationSource.get_or_create_from_instance(snippet)
+
+        source.update_target_view_restrictions(self.fr_locale)
+        self.assertEqual(get_translated_instance.call_count, 0)
+
+    def test_sync_view_restrictions_raises_exception_for_snippets(self):
+        snippet = TestSnippet.objects.create(field="Test content")
+        source, created = TranslationSource.get_or_create_from_instance(snippet)
+
+        with self.assertRaises(NoViewRestrictionsError):
+            source.sync_view_restrictions(snippet, snippet)

wagtail_localize/tests/test_update_translations.py

@@ -1,6 +1,9 @@
+from unittest import mock
+
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import Group, Permission
 from django.contrib.contenttypes.models import ContentType
+from django.core.exceptions import ValidationError
 from django.test import TestCase, override_settings
 from django.urls import reverse
 from wagtail import VERSION as WAGTAIL_VERSION
@@ -573,3 +576,40 @@ def test_post_update_page_translation_after_source_privacy_changed(self):
         self.assertTrue(translated_restriction.restriction_type, "groups")
         self.assertEqual(translated_restriction.groups.count(), 1)
         self.assertEqual(translated_restriction.groups.first(), test_group)
+
+        another_group = Group.objects.create(name="Another Group")
+        restriction.groups.set([another_group])
+        self.client.post(
+            reverse(
+                "wagtail_localize:update_translations",
+                args=[self.page_source.id],
+            ),
+        )
+        self.fr_blog_post.refresh_from_db()
+        translated_restriction = self.fr_blog_post.view_restrictions.first()
+        self.assertEqual(translated_restriction.groups.first(), another_group)
+
+    @mock.patch(
+        "wagtail_localize.models.TranslationSource.update_target_view_restrictions"
+    )
+    def test_post_update_page_translation_will_run_update_target_view_restrictions(
+        self, update_target_view_restrictions
+    ):
+        self.client.post(
+            reverse(
+                "wagtail_localize:update_translations",
+                args=[self.page_source.id],
+            ),
+        )
+        self.assertEqual(update_target_view_restrictions.call_count, 1)
+
+        # now let's try with an exception
+        update_target_view_restrictions.side_effect = ValidationError("oops")
+        self.client.post(
+            reverse(
+                "wagtail_localize:update_translations",
+                args=[self.page_source.id],
+            ),
+        )
+        # one call from the first post, and one from above
+        self.assertEqual(update_target_view_restrictions.call_count, 2)