From 9235932f00b3acc858b938d61471b142a10f285e Mon Sep 17 00:00:00 2001
From: sewagodimo <edneciamat@gmail.com>
Date: Thu, 10 Jan 2019 12:41:53 +0200
Subject: [PATCH] update exclude varient format and add variants to tests

---
 src/wagtail_personalisation/utils.py | 22 +++++++++++-----------
 tests/unit/test_utils.py             |  7 ++++++-
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/wagtail_personalisation/utils.py b/src/wagtail_personalisation/utils.py
index bdc3c31..3579911 100644
--- a/src/wagtail_personalisation/utils.py
+++ b/src/wagtail_personalisation/utils.py
@@ -98,18 +98,18 @@ def parse_tag(token, parser):
 def exclude_variants(pages):
     """Checks if page is not a variant
 
-    :param pages: List of pages to check
-    :type pages: list
-    :return: List of pages that aren't variants
-    :rtype: list
+    :param pages: List | Queryset of pages to check
+    :type pages: list or querset
+    :return: List|Queryset of pages that aren't variants
+    :rtype: list or queryset (depending on the param type)
     """
 
     for page in pages:
-        if not ((hasattr(page, 'personalisation_metadata') is False
-                 ) or (hasattr(page, 'personalisation_metadata') and page.personalisation_metadata is None
-                       ) or (hasattr(page, 'personalisation_metadata') and page.personalisation_metadata.is_canonical)):
-                    if (type(pages) == list):
-                        pages.remove(page)
-                    else:
-                        pages.exclude(pk=page.pk)
+        if hasattr(page, 'personalisation_metadata') is not False and \
+           page.personalisation_metadata is not None and \
+           page.personalisation_metadata.is_canonical is not True:
+                if (type(pages) == list):
+                    pages.remove(page)
+                else:
+                    pages.exclude(pk=page.pk)
     return pages
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 0f04056..1c4a6a0 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -69,5 +69,10 @@ def test_exclude_variants_with_pages_querysets():
         page = WagtailPage(path="/" + str(i), depth=0, url_path="/", title="Hoi " + str(i))
         page.save()
     pages = WagtailPage.objects.all().order_by('id')
+    # add variants
+    for page in pages:
+        page.personalisation_metadata = Metadata(is_canonical=False)
+    pages = WagtailPage.objects.all().order_by('id')
     result = exclude_variants(pages)
-    assert result == pages
+    assert type(result) == type(pages)
+    assert result != pages