Refactor ProjectConfig to add attribute ID to key mapping and implement retrieval methods; update test for cmab field population

Author: FarhanAnjum-opti

optimizely/entities.py

@@ -194,4 +194,3 @@ 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,7 +109,8 @@ class IntegrationDict(BaseEntity):
     key: str
     host: str
     publicKey: str
-    
+
+
 class CmabDict(BaseEntity):
     """Cmab dict from parsed datafile json."""
     attributeIds: list[str]

optimizely/project_config.py

@@ -94,7 +94,9 @@ def __init__(self, datafile: str | bytes, logger: Logger, error_handler: Any):
         self.attribute_key_map: dict[str, entities.Attribute] = self._generate_key_map(
             self.attributes, 'key', entities.Attribute
         )
-
+        self.attribute_id_to_key_map: dict[str, str] = {}
+        for attribute in self.attributes:
+            self.attribute_id_to_key_map[attribute['id']] = attribute['key']
         self.audience_id_map: dict[str, entities.Audience] = self._generate_key_map(
             self.audiences, 'id', entities.Audience
         )
@@ -510,6 +512,34 @@ def get_attribute_id(self, attribute_key: str) -> Optional[str]:
         self.error_handler.handle_error(exceptions.InvalidAttributeException(enums.Errors.INVALID_ATTRIBUTE))
         return None
 
+    def get_attribute_by_key(self, key: str) -> Optional[entities.Attribute]:
+        """ Get attribute for the provided attribute key.
+
+        Args:
+            key: Attribute key for which attribute is to be fetched.
+
+        Returns:
+            Attribute corresponding to the provided attribute key.
+        """
+        if key in self.attribute_key_map:
+            return self.attribute_key_map[key]
+        self.logger.error(f'Attribute with key:"{key}" is not in datafile.')
+        return None
+
+    def get_attribute_key_by_id(self, id: str) -> Optional[str]:
+        """ Get attribute key for the provided attribute id.
+
+        Args:
+            id: Attribute id for which attribute is to be fetched.
+
+        Returns:
+            Attribute key corresponding to the provided attribute id.
+        """
+        if id in self.attribute_id_to_key_map:
+            return self.attribute_id_to_key_map[id]
+        self.logger.error(f'Attribute with id:"{id}" is not in datafile.')
+        return None
+
     def get_feature_from_key(self, feature_key: str) -> Optional[entities.FeatureFlag]:
         """ Get feature for the provided feature key.
 

tests/test_config.py

@@ -155,7 +155,7 @@ def test_init(self):
         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. """
+        """ Test that the cmab field is populated correctly in experiments."""
 
         # Deep copy existing datafile and add cmab config to the first experiment
         config_dict = copy.deepcopy(self.config_dict_with_multiple_experiments)
@@ -167,11 +167,9 @@ def test_cmab_field_population(self):
 
         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. """