Add more tests for wagtail_hooks

tests/unit/test_wagtail_hooks.py

@@ -1,5 +1,6 @@
 import pytest
 
+from tests.factories.segment import SegmentFactory
 from wagtail_personalisation import adapters, wagtail_hooks
 
 
@@ -39,3 +40,23 @@ def test_serve_variation_with_variant_segmented(site, rf, segmented_page):
 
     result = wagtail_hooks.serve_variation(page, request, args, kwargs)
     assert result.status_code == 200
+
+
+@pytest.mark.django_db
+def test_page_listing_variant_buttons(site, rf, segmented_page):
+    page = segmented_page.personalisation_metadata.canonical_page
+
+    SegmentFactory(name='something')
+    result = wagtail_hooks.page_listing_variant_buttons(page, [])
+    items = list(result)
+    assert len(items) == 1
+
+
+@pytest.mark.django_db
+def test_page_listing_more_buttons(site, rf, segmented_page):
+    page = segmented_page.personalisation_metadata.canonical_page
+
+    SegmentFactory(name='something')
+    result = wagtail_hooks.page_listing_more_buttons(page, [])
+    items = list(result)
+    assert len(items) == 2