Add CmabDict type and update Experiment class to include cmab field

Author: FarhanAnjum-opti

optimizely/entities.py

@@ -22,7 +22,7 @@
 
 if TYPE_CHECKING:
     # prevent circular dependenacy by skipping import at runtime
-    from .helpers.types import ExperimentDict, TrafficAllocation, VariableDict, VariationDict
+    from .helpers.types import ExperimentDict, TrafficAllocation, VariableDict, VariationDict, CmabDict
 
 
 class BaseEntity:
@@ -84,6 +84,7 @@ def __init__(
         audienceConditions: Optional[Sequence[str | list[str]]] = None,
         groupId: Optional[str] = None,
         groupPolicy: Optional[str] = None,
+        cmab: Optional[CmabDict] = None,
         **kwargs: Any
     ):
         self.id = id
@@ -97,6 +98,7 @@ def __init__(
         self.layerId = layerId
         self.groupId = groupId
         self.groupPolicy = groupPolicy
+        self.cmab = cmab
 
     def get_audience_conditions_or_ids(self) -> Sequence[str | list[str]]:
         """ Returns audienceConditions if present, otherwise audienceIds. """
@@ -192,3 +194,4 @@ def __init__(self, key: str, host: Optional[str] = None, publicKey: Optional[str
         self.key = key
         self.host = host
         self.publicKey = publicKey
+        

optimizely/helpers/types.py

@@ -109,3 +109,8 @@ class IntegrationDict(BaseEntity):
     key: str
     host: str
     publicKey: str
+    
+class CmabDict(BaseEntity):
+    """Cmab dict from parsed datafile json."""
+    attributeIds: list[str]
+    trafficAllocation: int

tests/test_config.py

@@ -154,6 +154,25 @@ def test_init(self):
         self.assertEqual(expected_variation_key_map, self.project_config.variation_key_map)
         self.assertEqual(expected_variation_id_map, self.project_config.variation_id_map)
 
+    def test_cmab_field_population(self):
+        """ Test that the cmab field is populated correctly in experiments and attributes are parsed as attribute_ids. """
+
+        # Deep copy existing datafile and add cmab config to the first experiment
+        config_dict = copy.deepcopy(self.config_dict_with_multiple_experiments)
+        config_dict['experiments'][0]['cmab'] = {'attributeIds': ['808797688', '808797689'], 'trafficAllocation': 4000}
+        config_dict['experiments'][0]['trafficAllocation'] = []
+
+        opt_obj = optimizely.Optimizely(json.dumps(config_dict))
+        project_config = opt_obj.config_manager.get_config()
+
+        experiment = project_config.get_experiment_from_key('test_experiment')
+        self.assertEqual(experiment.cmab, {'attributeIds': ['808797688', '808797689'], 'trafficAllocation': 4000})
+        
+        experiment_2 = project_config.get_experiment_from_key('test_experiment_2')
+        self.assertIsNone(experiment_2.cmab)
+        
+        print(project_config.attributes)
+
     def test_init__with_v4_datafile(self):
         """ Test that on creating object, properties are initiated correctly for version 4 datafile. """