feat: extracted-resolve-trait-values

Author: Zaimwa9

flagsmith/analytics.py

@@ -193,7 +193,7 @@ def _handle_flush_result(
             response = future.result()
             response.raise_for_status()
         except Exception:
-            logger.warning("Failed to flush pipeline analytics, re-queuing events")
+            logger.warning("Failed to flush pipeline analytics, re-queuing events", exc_info=True)
             with self._lock:
                 self._buffer = events + self._buffer
                 self._buffer = self._buffer[: self._max_buffer]

flagsmith/flagsmith.py

@@ -19,6 +19,7 @@
     map_environment_document_to_context,
     map_environment_document_to_environment_updated_at,
     map_segment_results_to_identity_segments,
+    resolve_trait_values,
 )
 from flagsmith.models import DefaultFlag, Flags, Segment
 from flagsmith.offline_handlers import OfflineHandler
@@ -318,21 +319,10 @@ def track_event(
         self._pipeline_analytics_processor.record_custom_event(
             event_name=event_name,
             identity_identifier=identity_identifier,
-            traits=self._resolve_traits(traits),
+            traits=resolve_trait_values(traits),
             metadata=metadata,
         )
 
-    @staticmethod
-    def _resolve_traits(
-        traits: typing.Optional[TraitMapping],
-    ) -> typing.Optional[typing.Dict[str, typing.Any]]:
-        if not traits:
-            return None
-        return {
-            key: (val["value"] if isinstance(val, dict) else val)
-            for key, val in traits.items()
-        }
-
     def update_environment(self) -> None:
         try:
             environment_data = self._get_json_response(
@@ -414,7 +404,7 @@ def _get_identity_flags_from_document(
             default_flag_handler=self.default_flag_handler,
             pipeline_analytics_processor=self._pipeline_analytics_processor,
             identity_identifier=identifier,
-            traits=self._resolve_traits(traits),
+            traits=resolve_trait_values(traits),
         )
 
     def _get_environment_flags_from_api(self) -> Flags:
@@ -461,7 +451,7 @@ def _get_identity_flags_from_api(
                 default_flag_handler=self.default_flag_handler,
                 pipeline_analytics_processor=self._pipeline_analytics_processor,
                 identity_identifier=identifier,
-                traits=self._resolve_traits(traits),
+                traits=resolve_trait_values(traits),
             )
         except FlagsmithAPIError:
             if self.offline_handler:

flagsmith/mappers.py

@@ -27,6 +27,7 @@
     SegmentMetadata,
     StreamEvent,
     TraitConfig,
+    TraitMapping,
 )
 from flagsmith.utils.datetime import fromisoformat
 
@@ -75,31 +76,27 @@ def map_environment_document_to_environment_updated_at(
     return updated_at.astimezone(tz=timezone.utc)
 
 
+def resolve_trait_values(
+    traits: typing.Optional[TraitMapping],
+) -> typing.Optional[typing.Dict[str, typing.Any]]:
+    if not traits:
+        return None
+    return {
+        key: (val["value"] if isinstance(val, dict) else val)
+        for key, val in traits.items()
+    }
+
+
 def map_context_and_identity_data_to_context(
     context: SDKEvaluationContext,
     identifier: str,
-    traits: typing.Optional[
-        typing.Mapping[
-            str,
-            typing.Union[
-                ContextValue,
-                TraitConfig,
-            ],
-        ]
-    ],
+    traits: typing.Optional[TraitMapping] = None,
 ) -> SDKEvaluationContext:
     return {
         **context,
         "identity": {
             "identifier": identifier,
-            "traits": {
-                trait_key: (
-                    trait_value_or_config["value"]
-                    if isinstance(trait_value_or_config, dict)
-                    else trait_value_or_config
-                )
-                for trait_key, trait_value_or_config in (traits or {}).items()
-            },
+            "traits": resolve_trait_values(traits) or {},
         },
     }
 

tests/test_flagsmith.py

@@ -997,3 +997,30 @@ def test_get_identity_flags_passes_identity_and_traits(
         identity_identifier="user123",
         traits={"plan": "premium"},
     )
+
+
+@responses.activate()
+def test_get_identity_flags_resolves_trait_config_values(
+    mocker: MockerFixture, api_key: str, identities_json: str
+) -> None:
+    config = PipelineAnalyticsConfig(analytics_server_url="http://test/")
+    flagsmith = Flagsmith(environment_key=api_key, pipeline_analytics_config=config)
+
+    mock_record = mocker.patch.object(
+        flagsmith._pipeline_analytics_processor, "record_evaluation_event"
+    )
+
+    responses.add(method="POST", url=flagsmith.identities_url, body=identities_json)
+    flags = flagsmith.get_identity_flags(
+        "user123",
+        traits={"plan": {"value": "premium", "transient": True}},
+    )
+    flags.get_flag("some_feature")
+
+    mock_record.assert_called_once_with(
+        flag_key="some_feature",
+        enabled=True,
+        value="some-value",
+        identity_identifier="user123",
+        traits={"plan": "premium"},
+    )