From ad6b8dcea4f633b4ee53f3a6ec11b2850bb4db83 Mon Sep 17 00:00:00 2001
From: "Storm B. Heg" <storm@stormbase.digital>
Date: Thu, 20 Jul 2023 17:32:41 +0200
Subject: [PATCH] Add BS4SoupClient wrapper to ease html testing

This wrapper allows for easy checking of the html responses by Django's
Test client. CSS-like selectors can be used to check for the presence of
certain html attributes.

No longer shall we resort to regexes.
---
 wagtail/test/utils/wagtail_tests.py | 42 +++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/wagtail/test/utils/wagtail_tests.py b/wagtail/test/utils/wagtail_tests.py
index fc738f9a82f0..51499c4f4c41 100644
--- a/wagtail/test/utils/wagtail_tests.py
+++ b/wagtail/test/utils/wagtail_tests.py
@@ -1,12 +1,54 @@
 import warnings
 from contextlib import contextmanager
 
+from bs4 import BeautifulSoup
 from django import VERSION as DJANGO_VERSION
 from django.contrib.auth import get_user_model
+from django.test.client import Client as DjangoTestClient
 from django.test.testcases import assert_and_parse_html
+from django.utils.functional import SimpleLazyObject
+
+
+class BS4SoupClient(DjangoTestClient):
+    """Convenience wrapper around Django's Client and BeautifulSoup's SoupSieve
+
+    The intend of this wrapper is to make it easier to check the html of a response.
+
+    Usage example:
+
+        client = BS4SoupClient()
+        response = client.get("/good-vibes/")
+
+        # Select a HTML <span class="good-vibes> in the response
+        good_vibes = response.select("span.good-vibes")
+
+        # Assert the is only one element on the page that matches the selector
+        self.assertEqual(len(soup), 1)
+    """
+
+    def _get_soup(self, response):
+        content_type = response.headers.get("Content-Type", None)
+        if "text/html" in content_type:
+            return BeautifulSoup(response.content, "html5lib")
+        raise ValueError(
+            "Cannot use BeatifulSoup, the response content type is not HTML"
+        )
+
+    def request(self, **request):
+        response = super().request(**request)
+        # Make the soup available lazily. We will only incur the cost of instantiating
+        # the soup upon the initial access of this attribute.
+        response.soup = SimpleLazyObject(lambda: self._get_soup(response))
+
+        # Alias these BeautifulSoup functions the the response for convenience
+        response.select = lambda selector: response.soup.select(selector)
+        response.select_one = lambda selector: response.soup.select_one(selector)
+        return response
 
 
 class WagtailTestUtils:
+    client_class = BS4SoupClient
+
     @staticmethod
     def create_test_user():
         """