Moving variations to use objects (#24)

Author: aliabbasrizvi

optimizely/bucketer.py

@@ -77,14 +77,14 @@ def _find_bucket(self, user_id, parent_id, traffic_allocations):
     return None
 
   def bucket(self, experiment, user_id):
-    """ For a given experiment and bucketing ID determines ID of variation to be shown to user.
+    """ For a given experiment and bucketing ID determines variation to be shown to user.
 
     Args:
       experiment: Object representing the experiment for which user is to be bucketed.
       user_id: ID for user.
 
     Returns:
-      Variation ID for variation in which the visitor with ID user_id will be put in. None if no variation.
+      Variation in which user with ID user_id will be put in. None if no variation.
     """
 
     if not experiment:
@@ -94,11 +94,11 @@ def bucket(self, experiment, user_id):
     forced_variations = experiment.forcedVariations
     if forced_variations and user_id in forced_variations:
       variation_key = forced_variations.get(user_id)
-      variation_id = self.config.get_variation_id(experiment.key, variation_key)
-      if variation_id:
+      variation = self.config.get_variation_from_key(experiment.key, variation_key)
+      if variation:
         self.config.logger.log(enums.LogLevels.INFO,
                                'User "%s" is forced in variation "%s".' % (user_id, variation_key))
-      return variation_id
+      return variation
 
     # Determine if experiment is in a mutually exclusive group
     if experiment.groupPolicy in GROUP_POLICIES:
@@ -123,10 +123,10 @@ def bucket(self, experiment, user_id):
     # Bucket user if not in white-list and in group (if any)
     variation_id = self._find_bucket(user_id, experiment.id, experiment.trafficAllocation)
     if variation_id:
-      variation_key = self.config.get_variation_key_from_id(experiment.key, variation_id)
+      variation = self.config.get_variation_from_id(experiment.key, variation_id)
       self.config.logger.log(enums.LogLevels.INFO, 'User "%s" is in variation "%s" of experiment %s.' %
-                             (user_id, variation_key, experiment.key))
-      return variation_id
+                             (user_id, variation.key, experiment.key))
+      return variation
 
     self.config.logger.log(enums.LogLevels.INFO, 'User "%s" is in no variation.' % user_id)
     return None

optimizely/entities.py

@@ -53,3 +53,10 @@ def __init__(self, id, policy, experiments, trafficAllocation):
     self.policy = policy
     self.experiments = experiments
     self.trafficAllocation = trafficAllocation
+
+
+class Variation(BaseEntity):
+
+  def __init__(self, id, key):
+    self.id = id
+    self.key = key

optimizely/event_builder.py

@@ -163,10 +163,10 @@ def _add_experiment_variation_params(self, user_id, valid_experiments):
     """
 
     for experiment in valid_experiments:
-      variation_id = self.bucketer.bucket(experiment, user_id)
-      if variation_id:
+      variation = self.bucketer.bucket(experiment, user_id)
+      if variation:
         self.params[self.EXPERIMENT_PARAM_FORMAT.format(experiment_prefix=self.EventParams.EXPERIMENT_PREFIX,
-                                                        experiment_id=experiment.id)] = variation_id
+                                                        experiment_id=experiment.id)] = variation.id
 
   def _add_conversion_goal(self, event_key, event_value):
     """ Add conversion goal information to the event.
@@ -337,14 +337,14 @@ def _add_required_params_for_conversion(self, event_key, user_id, event_value, v
 
     self.params[self.EventParams.LAYER_STATES] = []
     for experiment in valid_experiments:
-      variation_id = self.bucketer.bucket(experiment, user_id)
-      if variation_id:
+      variation = self.bucketer.bucket(experiment, user_id)
+      if variation:
         self.params[self.EventParams.LAYER_STATES].append({
           self.EventParams.LAYER_ID: experiment.layerId,
           self.EventParams.ACTION_TRIGGERED: True,
           self.EventParams.DECISION: {
             self.EventParams.EXPERIMENT_ID: experiment.id,
-            self.EventParams.VARIATION_ID: variation_id,
+            self.EventParams.VARIATION_ID: variation.id,
             self.EventParams.IS_LAYER_HOLDBACK: False
           }
         })

optimizely/optimizely.py

@@ -111,21 +111,21 @@ def activate(self, experiment_key, user_id, attributes=None):
       self.logger.log(enums.LogLevels.INFO, 'Not activating user "%s".' % user_id)
       return None
 
-    variation_id = self.bucketer.bucket(experiment, user_id)
+    variation = self.bucketer.bucket(experiment, user_id)
 
-    if not variation_id:
+    if not variation:
       self.logger.log(enums.LogLevels.INFO, 'Not activating user "%s".' % user_id)
       return None
 
     # Create and dispatch impression event
-    impression_event = self.event_builder.create_impression_event(experiment, variation_id, user_id, attributes)
+    impression_event = self.event_builder.create_impression_event(experiment, variation.id, user_id, attributes)
     self.logger.log(enums.LogLevels.INFO, 'Activating user "%s" in experiment "%s".' % (user_id, experiment.key))
     self.logger.log(enums.LogLevels.DEBUG,
                     'Dispatching impression event to URL %s with params %s.' % (impression_event.url,
                                                                                 impression_event.params))
     self.event_dispatcher.dispatch_event(impression_event)
 
-    return self.config.get_variation_key_from_id(experiment.key, variation_id)
+    return variation.key
 
   def track(self, event_key, user_id,  attributes=None, event_value=None):
     """ Send conversion event to Optimizely.
@@ -185,5 +185,10 @@ def get_variation(self, experiment_key, user_id, attributes=None):
 
     if not self._validate_preconditions(experiment, user_id, attributes):
       return None
-    variation_id = self.bucketer.bucket(experiment, user_id)
-    return self.config.get_variation_key_from_id(experiment_key, variation_id)
+    variation = self.bucketer.bucket(experiment, user_id)
+
+    if variation:
+      return variation.key
+
+    return None
+

optimizely/project_config.py

@@ -55,15 +55,14 @@ def __init__(self, datafile, logger, error_handler):
     self.experiment_id_map = {}
     self.variation_key_map = {}
     self.variation_id_map = {}
-    for experiment_key in self.experiment_key_map.keys():
-      experiment = self.experiment_key_map.get(experiment_key)
+    for experiment in self.experiment_key_map.values():
       self.experiment_id_map[experiment.id] = experiment
-      self.variation_key_map[experiment_key] = self._generate_key_map(
-        self.experiment_key_map.get(experiment_key).variations, 'key'
-      )
-      self.variation_id_map[experiment_key] = self._generate_key_map(
-        self.experiment_key_map.get(experiment_key).variations, 'id'
+      self.variation_key_map[experiment.key] = self._generate_key_map_entity(
+        experiment.variations, 'key', entities.Variation
       )
+      self.variation_id_map[experiment.key] = {}
+      for variation in self.variation_key_map.get(experiment.key).values():
+        self.variation_id_map[experiment.key][variation.id] = variation
 
   @staticmethod
   def _generate_key_map(list, key):
@@ -224,51 +223,51 @@ def get_audience(self, audience_id):
     self.logger.log(enums.LogLevels.ERROR, 'Audience ID "%s" is not in datafile.' % audience_id)
     self.error_handler.handle_error(exceptions.InvalidAudienceException((enums.Errors.INVALID_AUDIENCE_ERROR)))
 
-  def get_variation_key_from_id(self, experiment_key, variation_id):
-    """ Get variation key given experiment key and variation ID.
+  def get_variation_from_key(self, experiment_key, variation_key):
+    """ Get variation given experiment and variation key.
 
     Args:
-      experiment_key: Key representing parent experiment of variation.
-      variation_id: ID of the variation.
+      experiment: Key representing parent experiment of variation.
+      variation_key: Key representing the variation.
 
     Returns
-      Variation key.
+      Object representing the variation.
     """
 
-    variation_map = self.variation_id_map.get(experiment_key)
+    variation_map = self.variation_key_map.get(experiment_key)
 
     if variation_map:
-      variation_obj = variation_map.get(variation_id)
-      if variation_obj:
-        return variation_obj['key']
+      variation = variation_map.get(variation_key)
+      if variation:
+        return variation
       else:
-        self.logger.log(enums.LogLevels.ERROR, 'Variation ID "%s" is not in datafile.' % variation_id)
+        self.logger.log(enums.LogLevels.ERROR, 'Variation key "%s" is not in datafile.' % variation_key)
         self.error_handler.handle_error(exceptions.InvalidVariationException(enums.Errors.INVALID_VARIATION_ERROR))
         return None
 
     self.logger.log(enums.LogLevels.ERROR, 'Experiment key "%s" is not in datafile.' % experiment_key)
     self.error_handler.handle_error(exceptions.InvalidExperimentException(enums.Errors.INVALID_EXPERIMENT_KEY_ERROR))
     return None
 
-  def get_variation_id(self, experiment_key, variation_key):
-    """ Get variation ID given the experiment and variation key.
+  def get_variation_from_id(self, experiment_key, variation_id):
+    """ Get variation given experiment and variation ID.
 
     Args:
-      experiment_key: Parent experiment for the variation.
-      variation_key: Variation for which the ID is to be determined.
+      experiment: Key representing parent experiment of variation.
+      variation_id: ID representing the variation.
 
-    Returns:
-      Variation ID corresponding to the variation.
+    Returns
+      Object representing the variation.
     """
 
-    variation_map = self.variation_key_map.get(experiment_key)
+    variation_map = self.variation_id_map.get(experiment_key)
 
     if variation_map:
-      variation_obj = variation_map.get(variation_key)
-      if variation_obj:
-        return variation_obj['id']
+      variation = variation_map.get(variation_id)
+      if variation:
+        return variation
       else:
-        self.logger.log(enums.LogLevels.ERROR, 'Variation key "%s" is not in datafile.' % variation_key)
+        self.logger.log(enums.LogLevels.ERROR, 'Variation ID "%s" is not in datafile.' % variation_id)
         self.error_handler.handle_error(exceptions.InvalidVariationException(enums.Errors.INVALID_VARIATION_ERROR))
         return None
 

tests/test_bucketing.py

@@ -4,6 +4,7 @@
 import random
 
 from optimizely import bucketer
+from optimizely import entities
 from optimizely import logger
 from optimizely import optimizely
 from optimizely.helpers import enums
@@ -24,15 +25,18 @@ def test_bucket(self):
     # Variation 1
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value',
                     return_value=42) as mock_generate_bucket_value:
-      self.assertEqual('111128', self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
-                                                      'test_user'))
+      self.assertEqual(entities.Variation('111128', 'control'),
+                       self.bucketer.bucket(
+                         self.project_config.get_experiment_from_key('test_experiment'), 'test_user'
+                       ))
     mock_generate_bucket_value.assert_called_once_with('test_user111127')
 
     # Variation 2
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value',
                     return_value=4242) as mock_generate_bucket_value:
-      self.assertEqual('111129', self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
-                                                      'test_user'))
+      self.assertEqual(entities.Variation('111129', 'variation'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
+                                            'test_user'))
     mock_generate_bucket_value.assert_called_once_with('test_user111127')
 
     # No matching variation
@@ -52,8 +56,8 @@ def test_bucket__user_in_forced_variation(self):
     """ Test that bucket returns variation ID for variation user is forced in. """
 
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value') as mock_generate_bucket_value:
-      self.assertEqual('111128', self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
-                                                      'user_1'))
+      self.assertEqual(entities.Variation('111128', 'control'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'), 'user_1'))
 
     # Confirm that bucket value generation did not happen
     self.assertEqual(0, mock_generate_bucket_value.call_count)
@@ -62,7 +66,7 @@ def test_bucket__user_in_forced_variation__invalid_variation_id(self):
     """ Test that bucket returns None when variation user is forced in is invalid. """
 
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value') as mock_generate_bucket_value, \
-        mock.patch('optimizely.project_config.ProjectConfig.get_variation_id',
+        mock.patch('optimizely.project_config.ProjectConfig.get_variation_from_key',
                    return_value=None) as mock_get_variation_id:
       self.assertIsNone(self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
                                              'user_1'))
@@ -77,8 +81,8 @@ def test_bucket__experiment_in_group(self):
     # In group, matching experiment and variation
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value',
                     side_effect=[42, 4242]) as mock_generate_bucket_value:
-      self.assertEqual('28902', self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_1'),
-                                                     'test_user'))
+      self.assertEqual(entities.Variation('28902', 'group_exp_1_variation'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_1'), 'test_user'))
 
     self.assertEqual([mock.call('test_user19228'), mock.call('test_user32222')],
                      mock_generate_bucket_value.call_args_list)
@@ -110,8 +114,8 @@ def test_bucket__experiment_in_group__user_in_forced_variation(self):
     """ Test that bucket returns variation ID for variation user is forced in. """
 
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value') as mock_generate_bucket_value:
-      self.assertEqual('28905', self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_2'),
-                                                     'user_1'))
+      self.assertEqual(entities.Variation('28905', 'group_exp_2_control'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_2'), 'user_1'))
 
     # Confirm that bucket value generation did not happen
     self.assertEqual(0, mock_generate_bucket_value.call_count)
@@ -151,8 +155,9 @@ def test_bucket(self):
     # Variation 1
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value', return_value=42),\
         mock.patch('optimizely.logger.SimpleLogger.log') as mock_logging:
-      self.assertEqual('111128', self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
-                                                      'test_user'))
+      self.assertEqual(entities.Variation('111128', 'control'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
+                                            'test_user'))
 
     self.assertEqual(2, mock_logging.call_count)
     self.assertEqual(mock.call(enums.LogLevels.DEBUG, 'Assigned bucket 42 to user "test_user".'),
@@ -165,8 +170,9 @@ def test_bucket(self):
     # Variation 2
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value', return_value=4242),\
         mock.patch('optimizely.logger.SimpleLogger.log') as mock_logging:
-      self.assertEqual('111129', self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
-                                                      'test_user'))
+      self.assertEqual(entities.Variation('111129', 'variation'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
+                                            'test_user'))
     self.assertEqual(2, mock_logging.call_count)
     self.assertEqual(mock.call(enums.LogLevels.DEBUG, 'Assigned bucket 4242 to user "test_user".'),
                      mock_logging.call_args_list[0])
@@ -191,8 +197,8 @@ def test_bucket__user_in_forced_variation(self):
 
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value'),\
         mock.patch('optimizely.logger.SimpleLogger.log') as mock_logging:
-      self.assertEqual('111128', self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'),
-                                                      'user_1'))
+      self.assertEqual(entities.Variation('111128', 'control'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('test_experiment'), 'user_1'))
 
     mock_logging.assert_called_with(enums.LogLevels.INFO, 'User "user_1" is forced in variation "control".')
 
@@ -203,8 +209,8 @@ def test_bucket__experiment_in_group(self):
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value',
                     side_effect=[42, 4242]),\
         mock.patch('optimizely.logger.SimpleLogger.log') as mock_logging:
-      self.assertEqual('28902', self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_1'),
-                                                     'test_user'))
+      self.assertEqual(entities.Variation('28902', 'group_exp_1_variation'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_1'), 'test_user'))
     self.assertEqual(4, mock_logging.call_count)
     self.assertEqual(mock.call(enums.LogLevels.DEBUG, 'Assigned bucket 42 to user "test_user".'),
                      mock_logging.call_args_list[0])
@@ -277,8 +283,8 @@ def test_bucket__experiment_in_group__user_in_forced_variation(self):
 
     with mock.patch('optimizely.bucketer.Bucketer._generate_bucket_value'),\
         mock.patch('optimizely.logger.SimpleLogger.log') as mock_logging:
-      self.assertEqual('28905', self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_2'),
-                                                     'user_1'))
+      self.assertEqual(entities.Variation('28905', 'group_exp_2_control'),
+                       self.bucketer.bucket(self.project_config.get_experiment_from_key('group_exp_2'), 'user_1'))
 
     # Confirm that bucket value generation did not happen
     mock_logging.assert_called_with(enums.LogLevels.INFO, 'User "user_1" is forced in variation "group_exp_2_control".')