restore override_traits behaviour

Author: khvn26

flag_engine/engine.py

@@ -106,12 +106,10 @@ def _get_identity_feature_states_dict(
         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}
-        )
+    if override_traits and (identity_context := context["identity"]):
+        identity_context["traits"] = {
+            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(