From 9c5ab4facf12071fc661a7b9653be30b30a4e698 Mon Sep 17 00:00:00 2001
From: Matt Westcott <matt@west.co.tt>
Date: Mon, 27 Feb 2017 18:42:15 +0000
Subject: [PATCH] Add assertTagInTemplateScript test helper function

This asserts that a tag exists within a <script type="text/template"> block in the given HTML,
matching as per assertTagInHTML rules.
---
 wagtail/tests/utils.py                  | 37 +++++++++++++++++++++++++
 wagtail/wagtailcore/tests/test_tests.py | 17 ++++++++++++
 2 files changed, 54 insertions(+)
diff --git a/wagtail/tests/utils.py b/wagtail/tests/utils.py
index 9ea4cfe689..da5bfcb3a3 100644
--- a/wagtail/tests/utils.py
+++ b/wagtail/tests/utils.py
@@ -144,6 +144,22 @@ class WagtailTestUtils(object):
 
         return count
 
+    def _tag_is_template_script(self, tag):
+        if tag.name != 'script':
+            return False
+        return any(attr == ('type', 'text/template') for attr in tag.attributes)
+
+    def _find_template_script_tags(self, haystack):
+        if not hasattr(haystack, 'name'):
+            return
+
+        if self._tag_is_template_script(haystack):
+            yield haystack
+        else:
+            for child in haystack.children:
+                for script_tag in self._find_template_script_tags(child):
+                    yield script_tag
+
     def assertTagInHTML(self, needle, haystack, count=None, msg_prefix=''):
         needle = assert_and_parse_html(self, needle, None, 'First argument is not valid HTML:')
         haystack = assert_and_parse_html(self, haystack, None, 'Second argument is not valid HTML:')
@@ -159,6 +175,27 @@ class WagtailTestUtils(object):
     def assertNotInHTML(self, needle, haystack, msg_prefix=''):
         self.assertInHTML(needle, haystack, count=0, msg_prefix=msg_prefix)
 
+    def assertTagInTemplateScript(self, needle, haystack, count=None, msg_prefix=''):
+        needle = assert_and_parse_html(self, needle, None, 'First argument is not valid HTML:')
+        haystack = assert_and_parse_html(self, haystack, None, 'Second argument is not valid HTML:')
+        real_count = 0
+
+        for script_tag in self._find_template_script_tags(haystack):
+            if script_tag.children:
+                self.assertEqual(len(script_tag.children), 1)
+                script_html = assert_and_parse_html(
+                    self, script_tag.children[0], None, 'Script tag content is not valid HTML:'
+                )
+                real_count += self._count_tag_occurrences(needle, script_html)
+
+        if count is not None:
+            self.assertEqual(
+                real_count, count,
+                msg_prefix + "Found %d instances of '%s' in template script (expected %d)" % (real_count, needle, count)
+            )
+        else:
+            self.assertTrue(real_count != 0, msg_prefix + "Couldn't find '%s' in template script" % needle)
+
 
 class WagtailPageTests(WagtailTestUtils, TestCase):
     """
diff --git a/wagtail/wagtailcore/tests/test_tests.py b/wagtail/wagtailcore/tests/test_tests.py
index a94ca13870..8756361de4 100644
--- a/wagtail/wagtailcore/tests/test_tests.py
+++ b/wagtail/wagtailcore/tests/test_tests.py
@@ -29,6 +29,23 @@ class TestAssertTagInHTML(TestCase, WagtailTestUtils):
         with self.assertRaises(AssertionError):
             self.assertTagInHTML('<li lang="en" class="important really" data-extra="boom">', haystack)
 
+    def test_assert_tag_in_template_script(self):
+        haystack = """<html>
+            <script type="text/template">
+                <p class="really important">first template block</p>
+            </script>
+            <script type="text/template">
+                <p class="really important">second template block</p>
+            </script>
+            <p class="normal">not in a script tag</p>
+        </html>"""
+
+        self.assertTagInTemplateScript('<p class="important really">', haystack)
+        self.assertTagInTemplateScript('<p class="important really">', haystack, count=2)
+
+        with self.assertRaises(AssertionError):
+            self.assertTagInTemplateScript('<p class="normal">', haystack)
+
 
 class TestWagtailPageTests(WagtailPageTests):
     def setUp(self):
