From 1cd3c0c09c949b7c04d0a54dd85ef52f3ca16d12 Mon Sep 17 00:00:00 2001
From: Amanda Lim-Cabuloy <amanda.lim@torchbox.com>
Date: Mon, 13 Jan 2025 20:49:11 +0800
Subject: [PATCH] Add tests for the DivisionMixin

---
 tbx/core/tests/test_division_mixin.py | 98 +++++++++++++++++++++++++++
 1 file changed, 98 insertions(+)
 create mode 100644 tbx/core/tests/test_division_mixin.py

diff --git a/tbx/core/tests/test_division_mixin.py b/tbx/core/tests/test_division_mixin.py
new file mode 100644
index 000000000..3615611af
--- /dev/null
+++ b/tbx/core/tests/test_division_mixin.py
@@ -0,0 +1,98 @@
+from tbx.core.factories import HomePageFactory
+from tbx.divisions.factories import DivisionPageFactory
+from tbx.services.factories import ServiceAreaPageFactory
+from wagtail.models import Site
+from wagtail.test.utils import WagtailPageTestCase
+
+
+class TestDivisionMixin(WagtailPageTestCase):
+    def setUp(self):
+        super().setUp()
+
+        # Set up the site & homepage.
+        site = Site.objects.get(is_default_site=True)
+        root = site.root_page.specific
+        self.home = HomePageFactory(parent=root)
+
+        site.root_page = self.home
+        site.save()
+
+        # Set up a services "index" page.
+        self.services = ServiceAreaPageFactory(title="Services", parent=self.home)
+        self.service_1 = ServiceAreaPageFactory(title="Service 1", parent=self.services)
+        self.service_2 = ServiceAreaPageFactory(
+            title="Service 2", parent=self.service_1
+        )
+
+        # Set up a division page.
+        self.division_1 = DivisionPageFactory(parent=self.home)
+        self.division_2 = DivisionPageFactory(title="Public sector", parent=self.home)
+
+    def test_division_selected(self):
+        """
+        For a page that has division selected,
+        get_division should return the selected page.
+        """
+        self.service_1.division = self.division_1
+        self.service_1.save()
+
+        service_3 = ServiceAreaPageFactory(
+            division=self.division_2,
+            parent=self.service_2,
+            title="Service 3",
+        )
+
+        self.assertEqual(self.service_1.get_division(), self.division_1)
+        self.assertEqual(service_3.get_division(), self.division_2)
+
+    def test_division_selected_on_ancestor(self):
+        """
+        For a page that does not have a division selected
+        but an ancestor page has a division selected,
+        get_division should return the ancestor's selected page.
+        """
+        self.service_1.division = self.division_1
+        self.service_1.save()
+
+        service_3 = ServiceAreaPageFactory(
+            parent=self.service_2,
+            title="Service 3",
+        )
+
+        self.assertEqual(self.service_1.get_division(), self.division_1)
+        self.assertEqual(self.service_2.get_division(), self.division_1)
+        self.assertEqual(service_3.get_division(), self.division_1)
+
+    def test_division_as_ancestor(self):
+        """
+        For a page that does not have a division selected
+        but an ancestor page is a DivisionPage,
+        get_division should return the ancestor DivisionPage.
+        """
+        service_a = ServiceAreaPageFactory(
+            title="Service A",
+            parent=self.division_1,
+        )
+        service_b = ServiceAreaPageFactory(
+            title="Service A",
+            parent=service_a,
+        )
+
+        self.assertEqual(service_a.get_division(), self.division_1)
+        self.assertEqual(service_b.get_division(), self.division_1)
+
+    def test_no_division(self):
+        """
+        For a page that does not have a division selected
+        and is not a descendant of a DivisionPage
+        nor has an ancestor with a selected division,
+        get_division should return None.
+        """
+        service_3 = ServiceAreaPageFactory(
+            parent=self.service_2,
+            title="Service 3",
+        )
+
+        self.assertIsNone(self.service_1.get_division())
+        self.assertIsNone(self.service_2.get_division())
+        self.assertIsNone(service_3.get_division())