Move visit count logic to the segment adapter

This makes the code also a bit simpler and updates the tests

Author: mvantellingen

src/wagtail_personalisation/adapters.py

@@ -168,3 +168,19 @@ def refresh(self):
         self.request.session['segments'] = new_segments
 
         self.update_visit_count()
+
+    def add_page_visit(self, page):
+        visit_count = self.request.session.setdefault('visit_count', [])
+        page_visits = [visit for visit in visit_count if visit['id'] == page.pk]
+
+        if page_visits:
+            for page_visit in page_visits:
+                page_visit['count'] += 1
+            self.request.session.modified = True
+        else:
+            visit_count.append({
+                'slug': page.slug,
+                'id': page.pk,
+                'path': self.request.path,
+                'count': 1,
+            })

src/wagtail_personalisation/wagtail_hooks.py

@@ -39,41 +39,7 @@ def set_visit_count(page, request, serve_args, serve_kwargs):
     :type request: django.http.HttpRequest
 
     """
-    # Create a new dict in the session if it's empty.
-    if 'visit_count' not in request.session:
-        request.session['visit_count'] = []
-
-    def create_new_counter(page, request):
-        """Create a new counter dict and place it in session storage.
-
-        :param page: The page being served
-        :type page: wagtail.wagtailcore.models.Page
-        :param request: The http request
-        :type request: django.http.HttpRequest
-
-        """
-        countdict = {
-            "slug": page.slug,
-            "id": page.pk,
-            "path": request.path,
-            "count": 1,
-        }
-        request.session['visit_count'].append(countdict)
-
-    if len(request.session['visit_count']) > 0:
-        for index, counter in enumerate(request.session['visit_count']):
-            if counter['id'] == page.pk:
-                # Counter already exists. Increase the count value by 1.
-                newcount = counter['count'] + 1
-                request.session['visit_count'][index]['count'] = newcount
-                request.session.modified = True
-            else:
-                # Counter doesn't exist.
-                # Create a new counter with count value 1.
-                create_new_counter(page, request)
-    else:
-        # No counters exist. Create a new counter with count value 1.
-        create_new_counter(page, request)
+    request.segment_adapter.add_page_visit(page)
 
 
 @hooks.register('before_serve_page')

tests/conftest.py

@@ -2,6 +2,8 @@
 
 import pytest
 from wagtail.wagtailcore.models import Page, Site
+from wagtail_factories import SiteFactory
+from tests.factories.page import PageFactory
 
 
 @pytest.fixture(scope='session')
@@ -10,3 +12,10 @@ def django_db_setup(django_db_setup, django_db_blocker):
         # Remove some initial data that is brought by the sandbox module
         Site.objects.all().delete()
         Page.objects.all().exclude(depth=1).delete()
+
+
+@pytest.fixture(scope='function')
+def site():
+    site = SiteFactory(is_default_site=True)
+    PageFactory(parent=site.root_page, slug='page-1')
+    PageFactory(parent=site.root_page, slug='page-2')

tests/unit/test_middleware.py

@@ -6,6 +6,7 @@
 from freezegun import freeze_time
 from wagtail_factories import SiteFactory
 
+from tests.factories.page import PageFactory
 from tests.factories.rule import (
     DayRuleFactory, DeviceRuleFactory, QueryRuleFactory, ReferralRuleFactory,
     TimeRuleFactory, VisitCountRuleFactory)
@@ -257,19 +258,21 @@ def test_match_any_segmenting(self, client):
 
 
 @pytest.mark.django_db
-class TestUserVisitCount(object):
-
-    def setup(self):
-        self.site = SiteFactory(is_default_site=True)
-
-        # TODO: Set up a bunch of pages for testing the visit count
-
-    def test_visit_count(self, client):
-        client.get('/')
-
-        assert any(item['path'] == '/' for item in client.session['visit_count'])
-
-    def test_no_visit_count(self, client):
-        client.get('/')
-
-        assert not any(item['path'] == '/doesntexist' for item in client.session['visit_count'])
+def test_visit_count(site, client):
+    response = client.get('/')
+    assert response.status_code == 200
+    visit_count = client.session['visit_count']
+    assert visit_count[0]['path'] == '/'
+    assert visit_count[0]['count'] == 1
+
+    response = client.get('/')
+    assert response.status_code == 200
+    visit_count = client.session['visit_count']
+    assert visit_count[0]['path'] == '/'
+    assert visit_count[0]['count'] == 2
+
+    response = client.get('/page-1/')
+    assert response.status_code == 200
+    visit_count = client.session['visit_count']
+    assert visit_count[0]['count'] == 2
+    assert visit_count[1]['count'] == 1