refactor: Migrate from is_context_in_segment to get_evaluation_result (#6896)

Author: khvn26

api/environments/identities/helpers.py

@@ -2,6 +2,7 @@
 
 from django.db import models
 from django.db.models import Prefetch, Q
+from flag_engine.engine import get_evaluation_result
 
 from environments.identities.managers import IdentityManager
 from environments.identities.traits.models import Trait
@@ -11,15 +12,7 @@
 from features.multivariate.models import MultivariateFeatureStateValue
 from features.versioning.versioning_service import get_environment_flags_list
 from segments.models import Segment
-from util.engine_models.context.mappers import (
-    is_context_in_segment,
-    map_environment_identity_to_context,
-)
-from util.mappers.engine import (
-    map_identity_to_engine,
-    map_segment_to_engine,
-    map_traits_to_engine,
-)
+from util.mappers.engine import map_environment_to_evaluation_context
 
 
 class Identity(models.Model):
@@ -154,7 +147,6 @@ def get_segments(
         :param overrides_only: only retrieve the segments which have a valid override in the environment
         :return: List of matching segments
         """
-        matching_segments = []
         db_traits = (
             self.identity_traits.all() if (traits is None and self.id) else traits or []
         )
@@ -164,29 +156,18 @@ def get_segments(
         else:
             all_segments = self.environment.project.get_segments_from_cache()
 
-        engine_identity = map_identity_to_engine(
-            self,
-            with_overrides=False,
-            with_traits=False,
+        segments_by_pk = {segment.pk: segment for segment in all_segments}
+        context = map_environment_to_evaluation_context(
+            identity=self,
+            environment=self.environment,
+            traits=db_traits,
+            segments=all_segments,
         )
-        engine_traits = map_traits_to_engine(db_traits)
-
-        for segment in all_segments:
-            engine_segment = map_segment_to_engine(segment)
-
-            context = map_environment_identity_to_context(
-                environment=self.environment,
-                identity=engine_identity,
-                override_traits=engine_traits,
-            )
-
-            if is_context_in_segment(
-                context=context,
-                segment=engine_segment,
-            ):
-                matching_segments.append(segment)
-
-        return matching_segments
+        result = get_evaluation_result(context)
+        return [
+            segments_by_pk[segment_result["metadata"]["pk"]]
+            for segment_result in result["segments"]
+        ]
 
     def get_all_user_traits(self):  # type: ignore[no-untyped-def]
         # this is pointless, we should probably replace all uses with the below code

api/environments/identities/models.py

@@ -25,6 +25,7 @@
     LifecycleModelMixin,
     hook,
 )
+from flag_engine.utils.hashing import get_hashed_percentage_for_object_ids
 from ordered_model.models import OrderedModelBase  # type: ignore[import-untyped]
 from simple_history.models import HistoricalRecords  # type: ignore[import-untyped]
 
@@ -49,9 +50,6 @@
     SoftDeleteExportableModel,
     abstract_base_auditable_model_factory,
 )
-from environments.identities.helpers import (
-    get_hashed_percentage_for_object_ids,
-)
 from features.constants import ENVIRONMENT, FEATURE_SEGMENT, IDENTITY
 from features.custom_lifecycle import CustomLifecycleModelMixin
 from features.feature_states.models import AbstractBaseFeatureValueModel
@@ -750,8 +748,8 @@ def get_multivariate_feature_state_value(
         # avoid further queries to the DB
         mv_options = list(self.multivariate_feature_state_values.all())
 
-        percentage_value = (
-            get_hashed_percentage_for_object_ids([self.id, identity_hash_key]) * 100
+        percentage_value = get_hashed_percentage_for_object_ids(
+            [self.id, identity_hash_key]
         )
 
         # Iterate over the mv options in order of id (so we get the same value each

api/features/models.py

@@ -1,19 +1,15 @@
 import typing
 
 from django.db.models import Q
+from flag_engine.engine import get_evaluation_result
 from rest_framework import serializers
 
 from features.serializers import FeatureStateSerializerFull
 from integrations.common.serializers import (
     BaseEnvironmentIntegrationModelSerializer,
 )
 from segments.models import Segment
-from util.engine_models.context.mappers import is_context_in_segment
-from util.mappers.engine import (
-    map_engine_identity_to_context,
-    map_identity_to_engine,
-    map_segment_to_engine,
-)
+from util.mappers.engine import map_environment_to_evaluation_context
 
 from .models import WebhookConfiguration
 
@@ -32,16 +28,14 @@ class Meta:
         fields = ("id", "name", "member")
 
     def get_member(self, obj: Segment) -> bool:
-        engine_identity = map_identity_to_engine(
-            self.context.get("identity"),  # type: ignore[arg-type]
-            with_overrides=False,
-        )
-        engine_segment = map_segment_to_engine(obj)
-        context = map_engine_identity_to_context(engine_identity)
-        return is_context_in_segment(
-            context=context,
-            segment=engine_segment,
+        identity = self.context["identity"]
+        context = map_environment_to_evaluation_context(
+            identity=identity,
+            environment=identity.environment,
+            segments=[obj],
         )
+        result = get_evaluation_result(context)
+        return bool(result["segments"])
 
 
 class IntegrationFeatureStateSerializer(FeatureStateSerializerFull):

api/integrations/webhook/serializers.py

@@ -0,0 +1,5 @@
+from typing_extensions import TypedDict
+
+
+class SegmentEngineMetadata(TypedDict):
+    pk: int

api/segments/types.py

@@ -31,9 +31,9 @@
 @pytest.mark.parametrize(
     "hashed_percentage, expected_mv_value",
     (
-        (variant_1_percentage_allocation / 100 - 0.01, variant_1_value),
-        (total_variance_percentage / 100 - 0.01, variant_2_value),
-        (total_variance_percentage / 100 + 0.01, control_value),
+        (variant_1_percentage_allocation - 1, variant_1_value),
+        (total_variance_percentage - 1, variant_2_value),
+        (total_variance_percentage + 1, control_value),
     ),
 )
 @mock.patch("features.models.get_hashed_percentage_for_object_ids")

api/tests/integration/environments/identities/test_integration_identities.py

@@ -23,9 +23,6 @@
     STRING,
 )
 from environments.identities import views
-from environments.identities.helpers import (
-    get_hashed_percentage_for_object_ids,
-)
 from environments.identities.models import Identity
 from environments.identities.traits.models import Trait
 from environments.models import Environment, EnvironmentAPIKey
@@ -646,12 +643,9 @@ def test_identities_endpoint_returns_default_value_if_rule_type_percentage_split
         segment=segment, type=SegmentRule.ALL_RULE
     )
 
-    identity_percentage_value = get_hashed_percentage_for_object_ids(
-        [segment.id, identity.id]
-    )
     Condition.objects.create(
         operator=PERCENTAGE_SPLIT,
-        value=int(identity_percentage_value / 2),
+        value=0,
         rule=segment_rule,
     )
     feature_segment = FeatureSegment.objects.create(

api/tests/unit/environments/identities/test_unit_identities_helpers.py

@@ -594,7 +594,7 @@ def test_feature_state_type_feature_segment(
     assert feature_state.type == FEATURE_SEGMENT
 
 
-@pytest.mark.parametrize("hashed_percentage", (0.0, 0.3, 0.5, 0.8, 0.999999))
+@pytest.mark.parametrize("hashed_percentage", (0.0, 30.0, 50.0, 80.0, 99.9999))
 @mock.patch("features.models.get_hashed_percentage_for_object_ids")
 def test_get_multivariate_value_returns_correct_value_when_we_pass_identity(  # type: ignore[no-untyped-def]
     mock_get_hashed_percentage,