Modify get_field_blocks to return children of streamblocks

cfpb · May 10, 2018 · 9664dce · 9664dce
1 parent ba6c337
commit 9664dce
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 3 deletions.
diff --git a/wagtailinventory/helpers.py b/wagtailinventory/helpers.py
@@ -1,6 +1,6 @@
 from itertools import chain
 
-from wagtail.wagtailcore.blocks import ListBlock, StructBlock
+from wagtail.wagtailcore.blocks import ListBlock, StreamBlock, StructBlock
 from wagtail.wagtailcore.fields import StreamField
 
 from wagtailinventory.models import PageBlock
@@ -34,7 +34,7 @@ def get_field_blocks(value):
             child_blocks = value.bound_blocks.values()
         else:
             child_blocks = [value.value]
-    elif isinstance(block, ListBlock):
+    elif isinstance(block, (ListBlock, StreamBlock)):
         child_blocks = value.value
     else:
         child_blocks = []

diff --git a/wagtailinventory/tests/test_helpers.py b/wagtailinventory/tests/test_helpers.py
@@ -4,7 +4,8 @@
 
 from wagtailinventory.helpers import get_page_blocks
 from wagtailinventory.tests.testapp.models import (
-    MultipleStreamFieldsPage, NoStreamFieldsPage, SingleStreamFieldPage
+    MultipleStreamFieldsPage, NestedStreamBlockPage, NoStreamFieldsPage,
+    SingleStreamFieldPage
 )
 
 
@@ -121,3 +122,27 @@ def test_multiple_streamfields(self):
                 'wagtailinventory.tests.testapp.blocks.Organism',
             ]
         )
+
+    def test_streamfield_with_child_stream_block(self):
+        page = self.make_page_with_streamfields(
+            NestedStreamBlockPage,
+            content=[
+                {
+                    'type': 'streamblock',
+                    'value': [
+                        {
+                            'type': 'text',
+                            'value': 'Test content',
+                        },
+                    ],
+                }
+            ]
+        )
+
+        self.assertEqual(
+            get_page_blocks(page),
+            [
+                'wagtail.wagtailcore.blocks.field_block.CharBlock',
+                'wagtail.wagtailcore.blocks.stream_block.StreamBlock',
+            ]
+        )
diff --git a/wagtailinventory/tests/testapp/migrations/0002_nested_stream_block_page.py b/wagtailinventory/tests/testapp/migrations/0002_nested_stream_block_page.py
@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+import django.db.models.deletion
+import wagtail.wagtailcore.blocks
+import wagtail.wagtailcore.fields
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('testapp', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='NestedStreamBlockPage',
+            fields=[
+                ('page_ptr', models.OneToOneField(auto_created=True, on_delete=django.db.models.deletion.CASCADE, parent_link=True, primary_key=True, serialize=False, to='wagtailcore.Page')),
+                ('content', wagtail.wagtailcore.fields.StreamField([(b'streamblock', wagtail.wagtailcore.blocks.StreamBlock([(b'text', wagtail.wagtailcore.blocks.CharBlock()), (b'atom', wagtail.wagtailcore.blocks.StructBlock([(b'title', wagtail.wagtailcore.blocks.CharBlock())]))]))])),
+            ],
+            options={
+                'abstract': False,
+            },
+            bases=('wagtailcore.page',),
+        ),
+    ]
diff --git a/wagtailinventory/tests/testapp/models.py b/wagtailinventory/tests/testapp/models.py
@@ -44,3 +44,16 @@ class MultipleStreamFieldsPage(Page):
         FieldPanel('first'),
         FieldPanel('second'),
     ]
+
+
+class NestedStreamBlockPage(Page):
+    content = StreamField([
+        ('streamblock', wagtail_blocks.StreamBlock([
+            ('text', wagtail_blocks.CharBlock()),
+            ('atom', blocks.Atom()),
+        ])),
+    ])
+
+    content_panels = [
+        FieldPanel('content'),
+    ]