Use SnippetViewSet.icon in SnippetChooserBlock

wagtail/snippets/blocks.py

@@ -5,6 +5,12 @@ from wagtail.coreutils import resolve_model_string
 
 
 class SnippetChooserBlock(ChooserBlock):
+    # Blocks are instantiated before models are loaded, so we can't set
+    # self.meta.icon in __init__. We need to override the default value
+    # some time after the model is loaded, so we mark it as a mutable
+    # attribute and set it using set_meta_options.
+    MUTABLE_META_ATTRIBUTES = ["icon"]
+
     def __init__(self, target_model, **kwargs):
         super().__init__(**kwargs)
         self._target_model = target_model
@@ -17,7 +23,9 @@ class SnippetChooserBlock(ChooserBlock):
     def widget(self):
         from wagtail.snippets.widgets import AdminSnippetChooser
 
-        return AdminSnippetChooser(self.target_model)
+        # Override the default icon with the icon for the target model
+        self.set_meta_options({"icon": self.target_model.snippet_viewset.icon})
+        return AdminSnippetChooser(self.target_model, icon=self.meta.icon)
 
     class Meta:
         icon = "snippet"

wagtail/snippets/tests/test_viewset.py

@@ -5,8 +5,11 @@ from django.urls import reverse
 
 from wagtail.admin.admin_url_finder import AdminURLFinder
 from wagtail.admin.panels import get_edit_handler
+from wagtail.blocks.field_block import FieldBlockAdapter
 from wagtail.coreutils import get_dummy_request
 from wagtail.models import Locale, Workflow, WorkflowContentType
+from wagtail.snippets.blocks import SnippetChooserBlock
+from wagtail.snippets.widgets import AdminSnippetChooser
 from wagtail.test.testapp.models import Advert, FullFeaturedSnippet, SnippetChooserModel
 from wagtail.test.utils import WagtailTestUtils
 
@@ -83,6 +86,28 @@ class TestCustomIcon(WagtailTestUtils, TestCase):
         self.assertContains(response, "icon icon-cog")
 
 
+class TestSnippetChooserBlockWithIcon(TestCase):
+    def test_adapt(self):
+        block = SnippetChooserBlock(FullFeaturedSnippet)
+
+        block.set_name("test_snippetchooserblock")
+        js_args = FieldBlockAdapter().js_args(block)
+
+        self.assertEqual(js_args[0], "test_snippetchooserblock")
+        self.assertIsInstance(js_args[1], AdminSnippetChooser)
+        self.assertEqual(js_args[1].model, FullFeaturedSnippet)
+        # It should use the icon defined in the FullFeaturedSnippetViewSet
+        self.assertEqual(js_args[2]["icon"], "cog")
+
+    def test_deconstruct(self):
+        block = SnippetChooserBlock(FullFeaturedSnippet, required=False)
+        path, args, kwargs = block.deconstruct()
+        self.assertEqual(path, "wagtail.snippets.blocks.SnippetChooserBlock")
+        self.assertEqual(args, (FullFeaturedSnippet,))
+        # It should not add any extra kwargs for the icon
+        self.assertEqual(kwargs, {"required": False})
+
+
 class TestSnippetChooserPanelWithIcon(WagtailTestUtils, TestCase):
     def setUp(self):
         self.user = self.login()