Add method for retrieving segments for an identity (#24)

Author: Matthew Elwell

flagsmith/flagsmith.py

@@ -7,12 +7,13 @@
 from flag_engine.environments.builders import build_environment_model
 from flag_engine.environments.models import EnvironmentModel
 from flag_engine.identities.models import IdentityModel, TraitModel
+from flag_engine.segments.evaluator import get_identity_segments
 from requests.adapters import HTTPAdapter
 from urllib3 import Retry
 
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.exceptions import FlagsmithAPIError, FlagsmithClientError
-from flagsmith.models import DefaultFlag, Flags
+from flagsmith.models import DefaultFlag, Flags, Segment
 from flagsmith.polling_manager import EnvironmentDataPollingManager
 from flagsmith.utils.identities import generate_identities_data
 
@@ -43,7 +44,7 @@ def __init__(
         custom_headers: typing.Dict[str, typing.Any] = None,
         request_timeout_seconds: int = None,
         enable_local_evaluation: bool = False,
-        environment_refresh_interval_seconds: int = 60,
+        environment_refresh_interval_seconds: typing.Union[int, float] = 60,
         retries: Retry = None,
         enable_analytics: bool = False,
         default_flag_handler: typing.Callable[[str], DefaultFlag] = None,
@@ -129,6 +130,29 @@ def get_identity_flags(
             return self._get_identity_flags_from_document(identifier, traits)
         return self._get_identity_flags_from_api(identifier, traits)
 
+    def get_identity_segments(
+        self, identifier: str, traits: typing.Dict[str, typing.Any] = None
+    ) -> typing.List[Segment]:
+        """
+        Get a list of segments that the given identity is in.
+
+        :param identifier: a unique identifier for the identity in the current
+            environment, e.g. email address, username, uuid
+        :param traits: a dictionary of traits to add / update on the identity in
+            Flagsmith, e.g. {"num_orders": 10}
+        :return: list of Segment objects that the identity is part of.
+        """
+
+        if not self._environment:
+            raise FlagsmithClientError(
+                "Local evaluation required to obtain identity segments."
+            )
+
+        traits = traits or {}
+        identity_model = self._build_identity_model(identifier, **traits)
+        segment_models = get_identity_segments(self._environment, identity_model)
+        return [Segment(id=sm.id, name=sm.name) for sm in segment_models]
+
     def update_environment(self):
         self._environment = self._get_environment_from_api()
 

flagsmith/models.py

@@ -140,3 +140,9 @@ def get_flag(self, feature_name: str) -> BaseFlag:
             self._analytics_processor.track_feature(flag.feature_id)
 
         return flag
+
+
+@dataclass
+class Segment:
+    id: int
+    name: str

tests/conftest.py

@@ -4,10 +4,12 @@
 import string
 
 import pytest
+import responses
 from flag_engine.environments.builders import build_environment_model
 
 from flagsmith import Flagsmith
 from flagsmith.analytics import AnalyticsProcessor
+from flagsmith.flagsmith import DEFAULT_API_URL
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
 
@@ -24,6 +26,11 @@ def api_key():
     return "".join(random.sample(string.ascii_letters, 20))
 
 
+@pytest.fixture(scope="session")
+def server_api_key():
+    return "ser.%s" % "".join(random.sample(string.ascii_letters, 20))
+
+
 @pytest.fixture()
 def flagsmith(api_key):
     return Flagsmith(environment_key=api_key)
@@ -35,6 +42,26 @@ def environment_json():
         yield f.read()
 
 
+@pytest.fixture()
+def local_eval_flagsmith(server_api_key, environment_json, mocker):
+    mock_session = mocker.MagicMock()
+    mocker.patch("flagsmith.flagsmith.requests.Session", return_value=mock_session)
+
+    mock_environment_document_response = mocker.MagicMock(status_code=200)
+    mock_environment_document_response.json.return_value = json.loads(environment_json)
+    mock_session.get.return_value = mock_environment_document_response
+
+    flagsmith = Flagsmith(
+        environment_key=server_api_key,
+        enable_local_evaluation=True,
+        environment_refresh_interval_seconds=0.1,
+    )
+
+    yield flagsmith
+
+    flagsmith.__del__()
+
+
 @pytest.fixture()
 def environment_model(environment_json):
     return build_environment_model(json.loads(environment_json))

tests/data/environment.json

@@ -11,7 +11,30 @@
     },
     "id": 1,
     "hide_disabled_flags": false,
-    "segments": []
+    "segments": [
+      {
+        "id": 1,
+        "name": "Test segment",
+        "rules": [
+          {
+            "type": "ALL",
+            "rules": [
+              {
+                "type": "ALL",
+                "rules": [],
+                "conditions": [
+                  {
+                    "operator": "EQUAL",
+                    "property_": "foo",
+                    "value": "bar"
+                  }
+                ]
+              }
+            ]
+          }
+        ]
+      }
+    ]
   },
   "segment_overrides": [],
   "id": 1,

tests/test_flagsmith.py

@@ -336,3 +336,27 @@ def default_flag_handler(feature_name: str) -> DefaultFlag:
 
     # Then
     assert flags.get_flag("some-feature") == default_flag
+
+
+def test_get_identity_segments_no_traits(local_eval_flagsmith, environment_model):
+    # Given
+    identifier = "identifier"
+
+    # When
+    segments = local_eval_flagsmith.get_identity_segments(identifier)
+
+    # Then
+    assert segments == []
+
+
+def test_get_identity_segments_with_valid_trait(local_eval_flagsmith, environment_model):
+    # Given
+    identifier = "identifier"
+    traits = {"foo": "bar"}  # obtained from data/environment.json
+
+    # When
+    segments = local_eval_flagsmith.get_identity_segments(identifier, traits)
+
+    # Then
+    assert len(segments) == 1
+    assert segments[0].name == "Test segment"  # obtained from data/environment.json