Creating object for groups (#23)

Author: aliabbasrizvi

optimizely/bucketer.py

@@ -5,7 +5,6 @@
   from .lib import pymmh3 as mmh3
 
 from .helpers import enums
-from . import exceptions
 
 MAX_TRAFFIC_VALUE = 10000
 UNSIGNED_MAX_32_BIT_VALUE = 0xFFFFFFFF
@@ -103,16 +102,12 @@ def bucket(self, experiment, user_id):
 
     # Determine if experiment is in a mutually exclusive group
     if experiment.groupPolicy in GROUP_POLICIES:
-      group_traffic_allocations = self.config.get_traffic_allocation(self.config.group_id_map, experiment.groupId)
+      group = self.config.get_group(experiment.groupId)
 
-      if not group_traffic_allocations:
-        self.config.logger.log(enums.LogLevels.ERROR, 'Group ID "%s" is not in datafile.' % experiment.groupId)
-        self.config.error_handler.handle_error(
-          exceptions.InvalidExperimentException(enums.Errors.INVALID_GROUP_ID_ERROR)
-        )
+      if not group:
         return None
 
-      user_experiment_id = self._find_bucket(user_id, experiment.groupId, group_traffic_allocations)
+      user_experiment_id = self._find_bucket(user_id, experiment.groupId, group.trafficAllocation)
       if not user_experiment_id:
         self.config.logger.log(enums.LogLevels.INFO, 'User "%s" is in no experiment.' % user_id)
         return None

optimizely/entities.py

@@ -44,3 +44,12 @@ def __init__(self, id, key, status, audienceIds, variations, forcedVariations,
     self.layerId = layerId
     self.groupId = groupId
     self.groupPolicy = groupPolicy
+
+
+class Group(BaseEntity):
+
+  def __init__(self, id, policy, experiments, trafficAllocation):
+    self.id = id
+    self.policy = policy
+    self.experiments = experiments
+    self.trafficAllocation = trafficAllocation

optimizely/exceptions.py

@@ -18,6 +18,11 @@ class InvalidEventException(Exception):
   pass
 
 
+class InvalidGroupException(Exception):
+  """ Raised when provided group ID is invalid. """
+  pass
+
+
 class InvalidVariationException(Exception):
   """ Raised when provided variation is invalid. """
   pass

optimizely/project_config.py

@@ -37,18 +37,18 @@ def __init__(self, datafile, logger, error_handler):
     self.audiences = config.get('audiences', [])
 
     # Utility maps for quick lookup
-    self.group_id_map = self._generate_key_map(self.groups, 'id')
+    self.group_id_map = self._generate_key_map_entity(self.groups, 'id', entities.Group)
     self.experiment_key_map = self._generate_key_map_entity(self.experiments, 'key', entities.Experiment)
     self.event_key_map = self._generate_key_map_entity(self.events, 'key', entities.Event)
     self.attribute_key_map = self._generate_key_map_entity(self.attributes, 'key', entities.Attribute)
     self.audience_id_map = self._generate_key_map_entity(self.audiences, 'id', entities.Audience)
     self.audience_id_map = self._deserialize_audience(self.audience_id_map)
     for group in self.group_id_map.values():
-      experiments_in_group_key_map = self._generate_key_map_entity(group['experiments'], 'key', entities.Experiment)
+      experiments_in_group_key_map = self._generate_key_map_entity(group.experiments, 'key', entities.Experiment)
       for experiment in experiments_in_group_key_map.values():
         experiment.__dict__.update({
-          'groupId': group.get('id'),
-          'groupPolicy': group.get('policy')
+          'groupId': group.id,
+          'groupPolicy': group.policy
         })
       self.experiment_key_map.update(experiments_in_group_key_map)
 
@@ -97,7 +97,6 @@ def _generate_key_map_entity(list, key, entity_class):
     """
 
     key_map = {}
-
     for obj in list:
       key_map[obj[key]] = entity_class(**obj)
 
@@ -188,6 +187,25 @@ def get_experiment_from_id(self, experiment_id):
     self.error_handler.handle_error(exceptions.InvalidExperimentException(enums.Errors.INVALID_EXPERIMENT_KEY_ERROR))
     return None
 
+  def get_group(self, group_id):
+    """ Get group for the provided group ID.
+
+    Args:
+      group_id: Group ID for which group is to be determined.
+
+    Returns:
+      Group corresponding to the provided group ID.
+    """
+
+    group = self.group_id_map.get(group_id)
+
+    if group:
+      return group
+
+    self.logger.log(enums.LogLevels.ERROR, 'Group ID "%s" is not in datafile.' % group_id)
+    self.error_handler.handle_error(exceptions.InvalidGroupException(enums.Errors.INVALID_GROUP_ID_ERROR))
+    return None
+
   def get_audience(self, audience_id):
     """ Get audience object for the provided audience ID.
 
@@ -304,21 +322,3 @@ 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 get_traffic_allocation(self, entity_key_map, entity_key):
-    """ Given an entity key map and entity key, returns the traffic allocation for that entity.
-
-    Args:
-      entity_key_map: Map representing the entity information.
-      entity_key: Key for whcih traffic allocation is to be retrieved from the map
-
-    Returns:
-      Traffic allocation for the experiment.
-    """
-
-    entity = entity_key_map.get(entity_key)
-
-    if entity:
-      return entity.get('trafficAllocation')
-
-    return None

tests/base.py

@@ -46,7 +46,6 @@ def setUp(self):
       'groups': [{
         'id': '19228',
         'policy': 'random',
-        'holdback': 0,
         'experiments': [{
           'id': '32222',
           'key': 'group_exp_1',
@@ -164,7 +163,6 @@ def setUp(self):
       'groups': [{
         'id': '19228',
         'policy': 'random',
-        'holdback': 0,
         'experiments': [{
           'id': '32222',
           'key': 'group_exp_1',

tests/test_config.py

@@ -23,7 +23,12 @@ def test_init(self):
     self.assertEqual(self.config_dict['experiments'], self.project_config.experiments)
     self.assertEqual(self.config_dict['events'], self.project_config.events)
     expected_group_id_map = {
-      '19228': copy.deepcopy(self.config_dict['groups'][0])
+      '19228': entities.Group(
+        self.config_dict['groups'][0]['id'],
+        self.config_dict['groups'][0]['policy'],
+        self.config_dict['groups'][0]['experiments'],
+        self.config_dict['groups'][0]['trafficAllocation']
+      )
     }
     expected_experiment_key_map = {
       'test_experiment': entities.Experiment(
@@ -316,18 +321,20 @@ def test_get_attribute__invalid_key(self):
 
     self.assertIsNone(self.project_config.get_attribute('invalid_key'))
 
-  def test_get_traffic_allocation__valid_key(self):
-    """ Test that trafficAllocation is retrieved correctly for valid group ID. """
+  def test_get_group__valid_id(self):
+    """ Test that group is retrieved correctly for valid group ID. """
+
+    self.assertEqual(entities.Group(self.config_dict['groups'][0]['id'],
+                                    self.config_dict['groups'][0]['policy'],
+                                    self.config_dict['groups'][0]['experiments'],
+                                    self.config_dict['groups'][0]['trafficAllocation']),
+                     self.project_config.get_group('19228'))
 
-    self.assertEqual(self.config_dict['groups'][0]['trafficAllocation'],
-                     self.project_config.get_traffic_allocation(self.project_config.group_id_map,
-                                                                '19228'))
 
-  def test_get_traffic_allocation__invalid_key(self):
+  def test_get_group__invalid_id(self):
     """ Test that None is returned when provided group ID is invalid. """
 
-    self.assertIsNone(self.project_config.get_traffic_allocation(self.project_config.group_id_map,
-                                                                 'invalid_key'))
+    self.assertIsNone(self.project_config.get_group('42'))
 
 
 class ConfigTestV2(base.BaseTestV2):
@@ -446,6 +453,14 @@ def test_get_attribute__invalid_key(self):
 
     mock_logging.assert_called_once_with(enums.LogLevels.ERROR, 'Attribute "invalid_key" is not in datafile.')
 
+  def test_get_group__invalid_id(self):
+    """ Test that message is logged when provided group ID is invalid. """
+
+    with mock.patch('optimizely.logger.SimpleLogger.log') as mock_logging:
+      self.project_config.get_group('42')
+
+    mock_logging.assert_called_once_with(enums.LogLevels.ERROR, 'Group ID "42" is not in datafile.')
+
 
 class ConfigExceptionTest(base.BaseTestV1):
 
@@ -503,3 +518,10 @@ def test_get_attribute__invalid_key(self):
     self.assertRaisesRegexp(exceptions.InvalidAttributeException,
                             enums.Errors.INVALID_ATTRIBUTE_ERROR,
                             self.project_config.get_attribute, 'invalid_key')
+
+  def test_get_group__invalid_id(self):
+    """ Test that exception is raised when provided group ID is invalid. """
+
+    self.assertRaisesRegexp(exceptions.InvalidGroupException,
+                            enums.Errors.INVALID_GROUP_ID_ERROR,
+                            self.project_config.get_group, '42')