Merge user entries on Login

This allows e.g. testing signup/login goals properly.

Author: theospears

experiments/__init__.py

@@ -1 +1,6 @@
+from django.contrib.auth.signals import user_logged_in
+from experiments.signals import transfer_enrollments_to_user
+
 from experiments.utils import _record_goal as record_goal
+
+user_logged_in.connect(transfer_enrollments_to_user, dispatch_uid="experiments_user_logged_in")

experiments/counters.py

@@ -13,20 +13,33 @@
 COUNTER_CACHE_KEY = 'experiments:participants:%s'
 COUNTER_FREQ_CACHE_KEY = 'experiments:freq:%s'
 
-def increment(key, participant_identifier):
+def increment(key, participant_identifier, count=1):
+    if count == 0:
+        return
+
     try:
         cache_key = COUNTER_CACHE_KEY % key
         freq_cache_key = COUNTER_FREQ_CACHE_KEY % key
-        new_value = r.hincrby(cache_key, participant_identifier, 1)
+        new_value = r.hincrby(cache_key, participant_identifier, count)
 
         # Maintain histogram of per-user counts
-        if new_value > 1:
-            r.hincrby(freq_cache_key, new_value - 1, -1)
+        if new_value > count:
+            r.hincrby(freq_cache_key, new_value - count, -1)
         r.hincrby(freq_cache_key, new_value, 1)
     except (ConnectionError, ResponseError):
         # Handle Redis failures gracefully
         pass
 
+def clear(key, participant_identifier):
+    # Remove the direct entry
+    cache_key = COUNTER_CACHE_KEY % key
+    pipe = r.pipeline()
+    freq, _ = pipe.hget(key, participant_identifier).hdel(cache_key, participant_identifier).execute()
+
+    # Remove from the histogram
+    freq_cache_key = COUNTER_FREQ_CACHE_KEY % key
+    r.hincrby(freq_cache_key, freq, -1)
+
 def get(key):
     try:
         cache_key = COUNTER_CACHE_KEY % key
@@ -35,6 +48,14 @@ def get(key):
         # Handle Redis failures gracefully
         return 0
 
+def get_frequency(key, participant_identifier):
+    try:
+        cache_key = COUNTER_CACHE_KEY % key
+        freq = r.hget(cache_key, participant_identifier)
+        return int(freq) if freq else 0
+    except (ConnectionError, ResponseError):
+        # Handle Redis failures gracefully
+        return 0
 
 def get_frequencies(key):
     try:

experiments/models.py

@@ -84,17 +84,33 @@ def increment_participant_count(self, alternative_name, participant_identifier):
         counter_key = PARTICIPANT_KEY % (self.name, alternative_name)
         counters.increment(counter_key, participant_identifier)
 
-    def increment_goal_count(self, alternative_name, goal_name, participant_identifier):
+    def increment_goal_count(self, alternative_name, goal_name, participant_identifier, count=1):
         # Increment experiment_name:alternative:participant counter
         counter_key = GOAL_KEY % (self.name, alternative_name, goal_name)
-        counters.increment(counter_key, participant_identifier)
+        counters.increment(counter_key, participant_identifier, count)
+
+    def remove_participant(self, alternative_name, participant_identifier):
+        # Remove participation record
+        counter_key = PARTICIPANT_KEY % (self.name, alternative_name)
+        counters.clear(counter_key, participant_identifier)
+
+        # Remove goal records
+        goal_names = getattr(settings, 'EXPERIMENTS_GOALS', [])
+        for goal_name in goal_names:
+            counter_key = GOAL_KEY % (self.name, alternative_name, goal_name)
+            counters.clear(counter_key, participant_identifier)
 
     def participant_count(self, alternative):
         return counters.get(PARTICIPANT_KEY % (self.name, alternative))
 
     def goal_count(self, alternative, goal):
         return counters.get(GOAL_KEY % (self.name, alternative, goal))
 
+    def participant_goal_frequencies(self, alternative, participant_identifier):
+        goal_names = getattr(settings, 'EXPERIMENTS_GOALS', [])
+        for goal in goal_names:
+            yield goal, counters.get_frequency(GOAL_KEY % (self.name, alternative, goal), participant_identifier)
+
     def goal_distribution(self, alternative, goal):
         return counters.get_frequencies(GOAL_KEY % (self.name, alternative, goal))
 

experiments/signals.py

@@ -0,0 +1,8 @@
+from __future__ import absolute_import
+
+from experiments.utils import create_user
+
+def transfer_enrollments_to_user(sender, request, user, **kwargs):
+    anon_user = create_user(session=request.session)
+    authenticated_user = create_user(user=user)
+    authenticated_user.incorporate(anon_user)

experiments/tests/__init__.py

@@ -4,3 +4,7 @@
 from .mannwhitney import MannWhitneyTestCase
 from .counter import CounterTestCase
 from .webuser import WebUserAnonymousTestCase, WebUserAuthenticatedTestCase, BotTestCase
+from . import webuser_incorporate
+
+def load_tests(*args, **kwargs):
+    return webuser_incorporate.load_tests(*args, **kwargs)

experiments/tests/webuser_incorporate.py

@@ -0,0 +1,60 @@
+from django.test import TestCase
+from django.utils.unittest import TestSuite
+from django.contrib.sessions.backends.db import SessionStore as DatabaseSession
+from django.contrib.auth.models import User
+
+from experiments.utils import DummyUser, SessionUser, AuthenticatedUser
+from experiments.models import Experiment, ENABLED_STATE
+
+class WebUserIncorporateTestCase(object):
+    def test_can_incorporate(self):
+        self.incorporating.incorporate(self.incorporated)
+
+    def test_incorporates_enrollment_from_other(self):
+        if not self._has_data():
+            return
+
+        try:
+            experiment = Experiment.objects.create(name='backgroundcolor', state=ENABLED_STATE)
+            self.incorporated.set_enrollment(experiment, 'blue')
+            self.incorporating.incorporate(self.incorporated)
+            self.assertEqual(self.incorporating.get_enrollment(experiment), 'blue')
+        finally:
+            experiment.delete()
+
+    def _has_data(self):
+        return not isinstance(self.incorporated, DummyUser) and not isinstance(self.incorporating, DummyUser)
+
+def dummy(incorporating):
+    return DummyUser()
+
+def anonymous(incorporating):
+    return SessionUser(session=DatabaseSession())
+
+def authenticated(incorporating):
+    return AuthenticatedUser(user=User.objects.create(username=['incorporating_user', 'incorporated_user'][incorporating]))
+
+user_factories = (dummy, anonymous, authenticated)
+
+def load_tests(loader, standard_tests, _):
+    suite = TestSuite()
+    suite.addTests(standard_tests)
+
+    for incorporating in user_factories:
+        for incorporated in user_factories:
+            test_case = build_test_case(incorporating, incorporated)
+            tests = loader.loadTestsFromTestCase(test_case)
+            suite.addTests(tests)
+    return suite
+
+
+def build_test_case(incorporating, incorporated):
+    class InstantiatedTestCase(WebUserIncorporateTestCase, TestCase):
+
+        def setUp(self):
+            super(InstantiatedTestCase, self).setUp()
+            self.incorporating = incorporating(True)
+            self.incorporated = incorporated(False)
+    InstantiatedTestCase.__name__ = "WebUserIncorporateTestCase_into_%s_from_%s" % (incorporating.__name__, incorporated.__name__)
+    return InstantiatedTestCase
+

experiments/utils.py

@@ -54,7 +54,7 @@ def set_enrollment(self, experiment, alternative):
         alternative will be increment, but those for the old one will not be decremented."""
         raise NotImplementedError
 
-    def record_goal(self, goal_name):
+    def record_goal(self, goal_name, count=1):
         """Record that this user has performed a particular goal
 
         This will update the goal stats for all experiments the user is enrolled in."""
@@ -85,16 +85,62 @@ def is_enrolled(self, experiment_name, alternative, request):
 
         return alternative == chosen_alternative
 
+    def incorporate(self, other_user):
+        """Incorporate all enrollments and goals performed by the other user
+
+        If this user is not enrolled in a given experiment, the results for the
+        other user are incorporated. For experiments this user is already
+        enrolled in the results of the other user are discarded.
+
+        This takes a relatively large amount of time for each experiment the other
+        user is enrolled in."""
+        for experiment, alternative in other_user._get_all_enrollments():
+            if not self._is_enrolled_in_experiment(experiment):
+                self.set_enrollment(experiment, alternative)
+                goals = experiment.participant_goal_frequencies(alternative, other_user._participant_identifier())
+                for goal_name, count in goals:
+                    experiment.increment_goal_count(alternative, goal_name, self._participant_identifier(), count)
+            other_user._cancel_enrollment(experiment)
+
+    def _participant_identifier(self):
+        "Unique identifier for this user in the counter store"
+        raise NotImplementedError
+
+    def _get_all_enrollments(self):
+        "Return experiment, alternative tuples for all experiments the user is enrolled in"
+        raise NotImplementedError
+
+    def _is_enrolled_in_experiment(self, experiment):
+        "Test whether the user currently has an enrollment in the supplied experiment"
+        raise NotImplementedError
+
+    def _cancel_enrollment(self, experiment):
+        "Remove the enrollment and any goals the user has against this experiment"
+        raise NotImplementedError
+
 
 class DummyUser(WebUser):
     def get_enrollment(self, experiment):
         return CONTROL_GROUP
     def set_enrollment(self, experiment, alternative):
         pass
-    def record_goal(self, goal_name):
+    def record_goal(self, goal_name, count=1):
         pass
     def is_enrolled(self, experiment_name, alternative, request):
         return alternative == CONTROL_GROUP
+    def incorporate(self, other_user):
+        for experiment, alternative in other_user._get_all_enrollments():
+            other_user._cancel_enrollment(experiment)
+    def _participant_identifier(self):
+        return ""
+    def _get_all_enrollments(self):
+        return []
+    def _is_enrolled_in_experiment(self, experiment):
+        return False
+    def _cancel_enrollment(self, experiment):
+        pass
+    def _get_goal_counts(self, experiment, alternative):
+        return {}
 
 
 class AuthenticatedUser(WebUser):
@@ -120,17 +166,32 @@ def set_enrollment(self, experiment, alternative):
             enrollment.save()
         experiment.increment_participant_count(alternative, self._participant_identifier())
 
-    def record_goal(self, goal_name):
-        enrollments = Enrollment.objects.filter(user=self.user)
-        if not enrollments:
-            return
-        for enrollment in enrollments: # Looks up by PK so no point caching.
-            if enrollment.experiment.is_displaying_alternatives():
-                enrollment.experiment.increment_goal_count(enrollment.alternative, goal_name, self._participant_identifier())
+    def record_goal(self, goal_name, count=1):
+        for experiment, alternative in self._get_all_enrollments():
+            if experiment.is_displaying_alternatives():
+                experiment.increment_goal_count(alternative, goal_name, self._participant_identifier(), count)
 
     def _participant_identifier(self):
         return 'user:%d' % (self.user.pk,)
 
+    def _get_all_enrollments(self):
+        enrollments = Enrollment.objects.filter(user=self.user).select_related("experiment")
+        if enrollments:
+            for enrollment in enrollments:
+                yield enrollment.experiment, enrollment.alternative
+
+    def _is_enrolled_in_experiment(self, experiment):
+        return Enrollment.objects.filter(user=self.user, experiment=experiment).exists()
+
+    def _cancel_enrollment(self, experiment):
+        try:
+            enrollment = Enrollment.objects.get(user=self.user, experiment=experiment)
+        except Enrollment.DoesNotExist:
+            pass
+        else:
+            experiment.remove_participant(enrollment.alternative, self._participant_identifier())
+            enrollment.delete()
+
 
 class SessionUser(WebUser):
     def __init__(self, session):
@@ -151,16 +212,11 @@ def set_enrollment(self, experiment, alternative):
         if self._is_verified_human():
             experiment.increment_participant_count(alternative, self._participant_identifier())
 
-    def record_goal(self, goal_name):
+    def record_goal(self, goal_name, count=1):
         if self._is_verified_human():
-            enrollments = self.session.get('experiments_enrollments', None)
-            if not enrollments:
-                return
-            for experiment_name, (alternative, goals) in enrollments.items():
-                experiment = experiment_manager.get(experiment_name, None)
-                if experiment and experiment.is_displaying_alternatives():
-                    experiment.increment_goal_count(alternative, goal_name, self._participant_identifier())
-            return
+            for experiment, alternative in self._get_all_enrollments():
+                if experiment.is_displaying_alternatives():
+                    experiment.increment_goal_count(alternative, goal_name, self._participant_identifier(), count)
         else:
             goals = self.session.get('experiments_goals', [])
             goals.append(goal_name) # Note, duplicates are allowed
@@ -172,16 +228,9 @@ def confirm_human(self):
 
         self.session['experiments_verified_human'] = True
 
-        enrollments = self.session.get('experiments_enrollments', None)
-        if not enrollments:
-            return
-
         # Replay enrollments
-        for experiment_name, data in enrollments.items():
-            alternative, goals = data
-            experiment = experiment_manager.get(experiment_name, None)
-            if experiment:
-                experiment.increment_participant_count(alternative, self._participant_identifier())
+        for experiment, alternative in self._get_all_enrollments():
+            experiment.increment_participant_count(alternative, self._participant_identifier())
 
         # Replay goals
         if 'experiments_goals' in self.session:
@@ -190,12 +239,34 @@ def confirm_human(self):
             del self.session['experiments_goals']
 
     def _participant_identifier(self):
-        if not self.session.session_key:
-            self.session.save() # Force session key
-        return 'session:%s' % (self.session.session_key,)
+        if 'experiments_session_key' not in self.session:
+            if not self.session.session_key:
+                self.session.save() # Force session key
+            self.session['experiments_session_key'] = self.session.session_key
+        return 'session:%s' % (self.session['experiments_session_key'],)
 
     def _is_verified_human(self):
         if getattr(settings, 'EXPERIMENTS_VERIFY_HUMAN', True):
             return self.session.get('experiments_verified_human', False)
         else:
             return True
+
+    def _get_all_enrollments(self):
+        enrollments = self.session.get('experiments_enrollments', None)
+        if enrollments:
+            for experiment_name, data in enrollments.items():
+                alternative, _ = data
+                experiment = experiment_manager.get(experiment_name, None)
+                if experiment:
+                    yield experiment, alternative
+
+    def _is_enrolled_in_experiment(self, experiment):
+        enrollments = self.session.get('experiments_enrollments', None)
+        return enrollments and experiment.name in enrollments
+
+    def _cancel_enrollment(self, experiment):
+        alternative = self.get_enrollment(experiment)
+        if alternative:
+            experiment.remove_participant(alternative, self._participant_identifier())
+            enrollments = self.session.get('experiments_enrollments', None)
+            del enrollments[experiment.name]