update: modify get_variation to return VariationResult and adjust related logic for improved variation handling

Author: FarhanAnjum-opti

optimizely/optimizely.py

@@ -29,7 +29,7 @@
 from .decision.optimizely_decide_option import OptimizelyDecideOption
 from .decision.optimizely_decision import OptimizelyDecision
 from .decision.optimizely_decision_message import OptimizelyDecisionMessage
-from .decision_service import Decision
+from .decision_service import Decision, VariationResult
 from .error_handler import NoOpErrorHandler, BaseErrorHandler
 from .event import event_factory, user_event_factory
 from .event.event_processor import BatchEventProcessor, BaseEventProcessor
@@ -535,8 +535,10 @@ def activate(self, experiment_key: str, user_id: str, attributes: Optional[UserA
             self.logger.error(enums.Errors.INVALID_PROJECT_CONFIG.format('activate'))
             return None
 
-        variation_key = self.get_variation(experiment_key, user_id, attributes)
-
+        variation_result = self.get_variation(experiment_key, user_id, attributes)
+        variation_key = None
+        if variation_result:
+            variation_key = variation_result['variation'].key
         if not variation_key:
             self.logger.info(f'Not activating user "{user_id}".')
             return None
@@ -612,17 +614,18 @@ def track(
 
     def get_variation(
         self, experiment_key: str, user_id: str, attributes: Optional[UserAttributes] = None
-    ) -> Optional[str]:
-        """ Gets variation where user will be bucketed.
+    ) -> Optional[VariationResult]:
+        """
+        Returns the variation result for the given user in the specified experiment.
 
         Args:
-          experiment_key: Experiment for which user variation needs to be determined.
-          user_id: ID for user.
-          attributes: Dict representing user attributes.
+            experiment_key: The key identifying the experiment.
+            user_id: The user ID.
+            attributes: Optional dictionary of user attributes.
 
         Returns:
-          Variation key representing the variation the user will be bucketed in.
-          None if user is not in experiment or if experiment is not Running.
+            A VariationResult object containing the variation assigned to the user, or None if the user is not
+            bucketed into any variation or the experiment is not running.
         """
 
         if not self.is_valid:
@@ -675,7 +678,7 @@ def get_variation(
             {'experiment_key': experiment_key, 'variation_key': variation_key},
         )
 
-        return variation_key
+        return variation_result
 
     def is_feature_enabled(self, feature_key: str, user_id: str, attributes: Optional[UserAttributes] = None) -> bool:
         """ Returns true if the feature is enabled for the given user.