From fdc0a7f2e1208d5a60cf96260c84e2cb8cf099d1 Mon Sep 17 00:00:00 2001
From: Kaitlyn Crawford <kaitlyn@praekelt.com>
Date: Wed, 21 Feb 2018 19:08:29 +0200
Subject: [PATCH] Get user info for Visit count rule

---
 src/wagtail_personalisation/rules.py | 22 ++++++++++++++++++++++
 tests/unit/test_rules_visitcount.py  | 21 +++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/src/wagtail_personalisation/rules.py b/src/wagtail_personalisation/rules.py
index b1bdb28..721e867 100644
--- a/src/wagtail_personalisation/rules.py
+++ b/src/wagtail_personalisation/rules.py
@@ -282,6 +282,28 @@ class VisitCountRule(AbstractBaseRule):
             ),
         }
 
+    def get_column_header(self):
+        return "Visit count - %s" % self.counted_page
+
+    def get_user_info_string(self, user):
+        # Local import for cyclic import
+        from wagtail_personalisation.adapters import (
+            get_segment_adapter, SessionSegmentsAdapter, SEGMENT_ADAPTER_CLASS)
+
+        # Create a fake request so we can use the adapter
+        request = RequestFactory().get('/')
+        request.user = user
+
+        # If we're using the session adapter check for an active session
+        if SEGMENT_ADAPTER_CLASS == SessionSegmentsAdapter:
+            request.session = self._get_user_session(user)
+        else:
+            request.session = SessionStore()
+
+        adapter = get_segment_adapter(request)
+        visit_count = adapter.get_visit_count(self.counted_page)
+        return str(visit_count)
+
 
 class QueryRule(AbstractBaseRule):
     """Query rule to segment users based on matching queries.
diff --git a/tests/unit/test_rules_visitcount.py b/tests/unit/test_rules_visitcount.py
index f153e3c..17dab9c 100644
--- a/tests/unit/test_rules_visitcount.py
+++ b/tests/unit/test_rules_visitcount.py
@@ -49,3 +49,24 @@ def test_visit_count_call_test_user_with_user_or_request_fails(site, client, use
     client.force_login(user)
 
     assert not rule.test_user(None)
+
+
+@pytest.mark.django_db
+def test_get_column_header(site):
+    segment = SegmentFactory(name='VisitCount')
+    rule = VisitCountRuleFactory(counted_page=site.root_page, segment=segment)
+
+    assert rule.get_column_header() == 'Visit count - Test page'
+
+
+@pytest.mark.django_db
+def test_get_user_info_string_returns_count(site, client, user):
+    segment = SegmentFactory(name='VisitCount')
+    rule = VisitCountRuleFactory(counted_page=site.root_page, segment=segment)
+
+    session = client.session
+    session['visit_count'] = [{'path': '/', 'count': 2}]
+    session.save()
+    client.force_login(user)
+
+    assert rule.get_user_info_string(user) == '2'