restore override_traits behaviour

Author: khvn26

flag_engine/context/mappers.py

@@ -1,11 +1,15 @@
+import typing
+
 from flag_engine.context.types import EvaluationContext
 from flag_engine.environments.models import EnvironmentModel
 from flag_engine.identities.models import IdentityModel
+from flag_engine.identities.traits.models import TraitModel
 
 
 def map_environment_identity_to_context(
     environment: EnvironmentModel,
     identity: IdentityModel,
+    override_traits: typing.Optional[typing.List[TraitModel]],
 ) -> EvaluationContext:
     """
     Maps an EnvironmentModel and IdentityModel to an EvaluationContext.
@@ -23,7 +27,12 @@ def map_environment_identity_to_context(
             "identifier": identity.identifier,
             "key": str(identity.django_id or identity.composite_key),
             "traits": {
-                trait.trait_key: trait.trait_value for trait in identity.identity_traits
+                trait.trait_key: trait.trait_value
+                for trait in (
+                    override_traits
+                    if override_traits is not None
+                    else identity.identity_traits
+                )
             },
         },
     }

flag_engine/engine.py

@@ -1,6 +1,7 @@
 import typing
 
 from flag_engine.context.mappers import map_environment_identity_to_context
+from flag_engine.context.types import EvaluationContext
 from flag_engine.environments.models import EnvironmentModel
 from flag_engine.features.models import FeatureModel, FeatureStateModel
 from flag_engine.identities.models import IdentityModel
@@ -54,9 +55,16 @@ def get_identity_feature_states(
     :return: list of feature state models based on the environment, any matching
         segments and any specific identity overrides
     """
+    context = map_environment_identity_to_context(
+        environment=environment,
+        identity=identity,
+        override_traits=override_traits,
+    )
+
     feature_states = list(
         _get_identity_feature_states_dict(
-            environment, identity, override_traits
+            environment=environment,
+            context=context,
         ).values()
     )
     if environment.get_hide_disabled_flags():
@@ -80,8 +88,15 @@ def get_identity_feature_state(
     :return: feature state model based on the environment, any matching
         segments and any specific identity overrides
     """
+    context = map_environment_identity_to_context(
+        environment=environment,
+        identity=identity,
+        override_traits=override_traits,
+    )
+
     feature_states = _get_identity_feature_states_dict(
-        environment, identity, override_traits
+        environment=environment,
+        context=context,
     )
     matching_feature = next(
         filter(lambda feature: feature.name == feature_name, feature_states.keys()),
@@ -96,23 +111,11 @@ def get_identity_feature_state(
 
 def _get_identity_feature_states_dict(
     environment: EnvironmentModel,
-    identity: IdentityModel,
-    override_traits: typing.Optional[typing.List[TraitModel]],
+    context: EvaluationContext,
 ) -> typing.Dict[FeatureModel, FeatureStateModel]:
     # Get feature states from the environment
     feature_states_by_feature = {fs.feature: fs for fs in environment.feature_states}
 
-    context = map_environment_identity_to_context(
-        environment=environment,
-        identity=identity,
-    )
-    if override_traits:
-        if typing.TYPE_CHECKING:  # pragma: no cover
-            assert context["identity"]
-        context["identity"].setdefault("traits", {}).update(
-            {trait.trait_key: trait.trait_value for trait in override_traits}
-        )
-
     # Override with any feature states defined by matching segments
     for context_segment in get_context_segments(
         context=context,