Alda/forced bucketing (#71)

* First pass at the forced bucketing feature.

* First pass the forced bucketing unit tests.

* Fixed the unit tests.

* Fixed lint errors.

* Responded to PR feeback.

* Responded to PR feeback. This file got left out from the previous commit.

* Code reduction per Ali.

* Fixed all of Ali's nits in the PR.

* Fixed lint error.

Author: alda-optimizely

optimizely/decision_service.py

@@ -99,6 +99,11 @@ def get_variation(self, experiment, user_id, attributes):
       self.logger.log(enums.LogLevels.INFO, 'Experiment "%s" is not running.' % experiment.key)
       return None
 
+    # Check if the user is forced into a variation
+    variation = self.config.get_forced_variation(experiment.key, user_id)
+    if variation:
+      return variation
+
     # Check to see if user is white-listed for a certain variation
     variation = self.get_forced_variation(experiment, user_id)
     if variation:

optimizely/optimizely.py

@@ -92,16 +92,16 @@ def _validate_instantiation_options(self, datafile, skip_json_validation):
     """
 
     if not skip_json_validation and not validator.is_datafile_valid(datafile):
-     raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('datafile'))
+      raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('datafile'))
 
     if not validator.is_event_dispatcher_valid(self.event_dispatcher):
-     raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('event_dispatcher'))
+      raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('event_dispatcher'))
 
     if not validator.is_logger_valid(self.logger):
-     raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('logger'))
+      raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('logger'))
 
     if not validator.is_error_handler_valid(self.error_handler):
-     raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('error_handler'))
+      raise exceptions.InvalidInputException(enums.Errors.INVALID_INPUT_ERROR.format('error_handler'))
 
   def _validate_user_inputs(self, attributes=None, event_tags=None):
     """ Helper method to validate user inputs.
@@ -261,6 +261,7 @@ def get_variation(self, experiment_key, user_id, attributes=None):
       return None
 
     experiment = self.config.get_experiment_from_key(experiment_key)
+
     if not experiment:
       self.logger.log(enums.LogLevels.INFO,
                       'Experiment key "%s" is invalid. Not activating user "%s".' % (experiment_key,
@@ -275,3 +276,33 @@ def get_variation(self, experiment_key, user_id, attributes=None):
       return variation.key
 
     return None
+
+  def set_forced_variation(self, experiment_key, user_id, variation_key):
+    """ Force a user into a variation for a given experiment.
+
+    Args:
+     experiment_key: A string key identifying the experiment.
+     user_id: The user ID.
+     variation_key: A string variation key that specifies the variation which the user.
+     will be forced into. If null, then clear the existing experiment-to-variation mapping.
+
+    Returns:
+      A boolean value that indicates if the set completed successfully.
+    """
+
+    return self.config.set_forced_variation(experiment_key, user_id, variation_key)
+
+  def get_forced_variation(self, experiment_key, user_id):
+    """ Gets the forced variation for a given user and experiment.
+
+      Args:
+       experiment_key: A string key identifying the experiment.
+       user_id: The user ID.
+
+     Returns:
+       The forced variation key. None if no forced variation key.
+    """
+
+    forced_variation = self.config.get_forced_variation(experiment_key, user_id)
+    return forced_variation.key if forced_variation else None
+

optimizely/project_config.py

@@ -84,6 +84,12 @@ def __init__(self, datafile, logger, error_handler):
 
     self.parsing_succeeded = True
 
+    # Map of user IDs to another map of experiments to variations.
+    # This contains all the forced variations set by the user
+    # by calling set_forced_variation (it is not the same as the
+    # whitelisting forcedVariations data structure).
+    self.forced_variation_map = {}
+
   @staticmethod
   def _generate_key_map(list, key, entity_class):
     """ Helper method to generate map from key to entity object for given list of dicts.
@@ -340,3 +346,107 @@ def get_attribute(self, attribute_key):
     self.logger.log(enums.LogLevels.ERROR, 'Attribute "%s" is not in datafile.' % attribute_key)
     self.error_handler.handle_error(exceptions.InvalidAttributeException(enums.Errors.INVALID_ATTRIBUTE_ERROR))
     return None
+
+  def set_forced_variation(self, experiment_key, user_id, variation_key):
+    """ Sets users to a map of experiments to forced variations.
+
+      Args:
+        experiment_key: Key for experiment.
+        user_id: The user ID.
+        variation_key: Key for variation. If None, then clear the existing experiment-to-variation mapping.
+
+      Returns:
+        A boolean value that indicates if the set completed successfully.
+    """
+    if not user_id:
+      self.logger.log(enums.LogLevels.DEBUG, 'User ID is invalid.')
+      return False
+
+    experiment = self.get_experiment_from_key(experiment_key)
+    if not experiment:
+      # The invalid experiment key will be logged inside this call.
+      return False
+
+    experiment_id = experiment.id
+    if not variation_key:
+      if user_id in self.forced_variation_map:
+        experiment_to_variation_map = self.forced_variation_map.get(user_id)
+        if experiment_id in experiment_to_variation_map:
+          del(self.forced_variation_map[user_id][experiment_id])
+          self.logger.log(enums.LogLevels.DEBUG,
+                          'Variation mapped to experiment "%s" has been removed for user "%s".'
+                          % (experiment_key, user_id))
+        else:
+          self.logger.log(enums.LogLevels.DEBUG,
+                          'Nothing to remove. Variation mapped to experiment "%s" for user "%s" does not exist.'
+                          % (experiment_key, user_id))
+      else:
+        self.logger.log(enums.LogLevels.DEBUG,
+                        'Nothing to remove. User "%s" does not exist in the forced variation map.' % user_id)
+      return True
+
+    forced_variation = self.get_variation_from_key(experiment_key, variation_key)
+    if not forced_variation:
+      # The invalid variation key will be logged inside this call.
+      return False
+
+    variation_id = forced_variation.id
+
+    if user_id not in self.forced_variation_map:
+      self.forced_variation_map[user_id] = {experiment_id: variation_id}
+    else:
+      self.forced_variation_map[user_id][experiment_id] = variation_id
+
+    self.logger.log(enums.LogLevels.DEBUG,
+                    'Set variation "%s" for experiment "%s" and user "%s" in the forced variation map.'
+                    % (variation_id, experiment_id, user_id))
+    return True
+
+  def get_forced_variation(self, experiment_key, user_id):
+    """ Gets the forced variation key for the given user and experiment.
+
+      Args:
+        experiment_key: Key for experiment.
+        user_id: The user ID.
+
+      Returns:
+        The variation which the given user and experiment should be forced into.
+    """
+    if not user_id:
+      self.logger.log(enums.LogLevels.DEBUG, 'User ID is invalid.')
+      return None
+
+    if user_id not in self.forced_variation_map:
+      self.logger.log(enums.LogLevels.DEBUG, 'User "%s" is not in the forced variation map.' % user_id)
+      return None
+
+    experiment = self.get_experiment_from_key(experiment_key)
+    if not experiment:
+      # The invalid experiment key will be logged inside this call.
+      return None
+
+    experiment_to_variation_map = self.forced_variation_map.get(user_id)
+
+    if not experiment_to_variation_map:
+      self.logger.log(enums.LogLevels.DEBUG,
+                      'No experiment "%s" mapped to user "%s" in the forced variation map.'
+                      % (experiment_key, user_id))
+      return None
+
+    variation_id = experiment_to_variation_map.get(experiment.id)
+    if variation_id is None:
+      self.logger.log(enums.LogLevels.DEBUG,
+                      'No variation mapped to experiment "%s" in the forced variation map.'
+                      % experiment_key)
+      return None
+
+    variation = self.get_variation_from_id(experiment_key, variation_id)
+    if not variation:
+      # The invalid variation ID will be logged inside this call.
+      return None
+
+    self.logger.log(enums.LogLevels.DEBUG,
+                    'Variation "%s" is mapped to experiment "%s" and user "%s" in the forced variation map'
+                    % (variation.key, experiment_key, user_id))
+    return variation
+

tests/test_config.py

@@ -459,6 +459,68 @@ def test_get_group__invalid_id(self):
 
     self.assertIsNone(self.project_config.get_group('42'))
 
+  # get_forced_variation tests
+  def test_get_forced_variation__invalid_user_id(self):
+    """ Test invalid user IDs return a null variation. """
+    self.project_config.forced_variation_map['test_user'] = {}
+    self.project_config.forced_variation_map['test_user']['test_experiment'] = 'test_variation'
+
+    self.assertIsNone(self.project_config.get_forced_variation('test_experiment', None))
+    self.assertIsNone(self.project_config.get_forced_variation('test_experiment', ''))
+
+  def test_get_forced_variation__invalid_experiment_key(self):
+    """ Test invalid experiment keys return a null variation. """
+    self.project_config.forced_variation_map['test_user'] = {}
+    self.project_config.forced_variation_map['test_user']['test_experiment'] = 'test_variation'
+
+    self.assertIsNone(self.project_config.get_forced_variation('test_experiment', None))
+    self.assertIsNone(self.project_config.get_forced_variation('test_experiment', ''))
+
+  # set_forced_variation tests
+  def test_set_forced_variation__invalid_user_id(self):
+    """ Test invalid user IDs set fail to set a forced variation """
+
+    self.assertFalse(self.project_config.set_forced_variation('test_experiment', None, 'variation'))
+    self.assertFalse(self.project_config.set_forced_variation('test_experiment', '', 'variation'))
+
+  def test_set_forced_variation__invalid_experiment_key(self):
+    """ Test invalid experiment keys set fail to set a forced variation """
+
+    self.assertFalse(self.project_config.set_forced_variation('test_experiment_not_in_datafile',
+                                                              'test_user', 'variation'))
+    self.assertFalse(self.project_config.set_forced_variation('', 'test_user', 'variation'))
+    self.assertFalse(self.project_config.set_forced_variation(None, 'test_user', 'variation'))
+
+  def test_set_forced_variation__invalid_variation_key(self):
+    """ Test invalid variation keys set fail to set a forced variation """
+
+    self.assertFalse(self.project_config.set_forced_variation('test_experiment', 'test_user',
+                                                              'variation_not_in_datafile'))
+    self.assertTrue(self.project_config.set_forced_variation('test_experiment', 'test_user', ''))
+    self.assertTrue(self.project_config.set_forced_variation('test_experiment', 'test_user', None))
+
+  def test_set_forced_variation__multiple_sets(self):
+    """ Test multiple sets of experiments for one and multiple users work """
+
+    self.assertTrue(self.project_config.set_forced_variation('test_experiment', 'test_user_1', 'variation'))
+    self.assertEqual(self.project_config.get_forced_variation('test_experiment', 'test_user_1').key, 'variation')
+    # same user, same experiment, different variation
+    self.assertTrue(self.project_config.set_forced_variation('test_experiment', 'test_user_1', 'control'))
+    self.assertEqual(self.project_config.get_forced_variation('test_experiment', 'test_user_1').key, 'control')
+    # same user, different experiment
+    self.assertTrue(self.project_config.set_forced_variation('group_exp_1', 'test_user_1', 'group_exp_1_control'))
+    self.assertEqual(self.project_config.get_forced_variation('group_exp_1', 'test_user_1').key, 'group_exp_1_control')
+
+    # different user
+    self.assertTrue(self.project_config.set_forced_variation('test_experiment', 'test_user_2', 'variation'))
+    self.assertEqual(self.project_config.get_forced_variation('test_experiment', 'test_user_2').key, 'variation')
+    # different user, different experiment
+    self.assertTrue(self.project_config.set_forced_variation('group_exp_1', 'test_user_2', 'group_exp_1_control'))
+    self.assertEqual(self.project_config.get_forced_variation('group_exp_1', 'test_user_2').key, 'group_exp_1_control')
+
+    # make sure the first user forced variations are still valid
+    self.assertEqual(self.project_config.get_forced_variation('test_experiment', 'test_user_1').key, 'control')
+    self.assertEqual(self.project_config.get_forced_variation('group_exp_1', 'test_user_1').key, 'group_exp_1_control')
 
 class ConfigLoggingTest(base.BaseTest):
   def setUp(self):