feat: Add indentity overrides to local evaluation mode (#72)

- Bump flag-engine
- Add environment overrides parsing
- Add pytest-cov
- Fix black pre-commit hook
- Fix Flag dataclasses

Author: khvn26

.gitignore

@@ -14,3 +14,5 @@ flagsmith.egg-info/
 
 .envrc
 .tool-versions
+
+.coverage

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ repos:
     hooks:
       - id: isort
   - repo: https://github.com/psf/black
-    rev: stable
+    rev: 23.3.0
     hooks:
       - id: black
         language_version: python3

flagsmith/flagsmith.py

@@ -7,7 +7,9 @@
 import requests
 from flag_engine import engine
 from flag_engine.environments.models import EnvironmentModel
-from flag_engine.identities.models import IdentityModel, TraitModel
+from flag_engine.identities.models import IdentityModel
+from flag_engine.identities.traits.models import TraitModel
+from flag_engine.identities.traits.types import TraitValue
 from flag_engine.segments.evaluator import get_identity_segments
 from requests.adapters import HTTPAdapter
 from urllib3 import Retry
@@ -94,6 +96,7 @@ def __init__(
         self.enable_realtime_updates = enable_realtime_updates
         self._analytics_processor = None
         self._environment = None
+        self._identity_overrides_by_identifier: typing.Dict[str, IdentityModel] = {}
 
         # argument validation
         if offline_mode and not offline_handler:
@@ -248,12 +251,21 @@ def get_identity_segments(
             )
 
         traits = traits or {}
-        identity_model = self._build_identity_model(identifier, **traits)
+        identity_model = self._get_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()
+        self._update_overrides()
+
+    def _update_overrides(self) -> None:
+        if not self._environment:
+            return
+        if overrides := self._environment.identity_overrides:
+            self._identity_overrides_by_identifier = {
+                identity.identifier: identity for identity in overrides
+            }
 
     def _get_environment_from_api(self) -> EnvironmentModel:
         environment_data = self._get_json_response(self.environment_url, method="GET")
@@ -269,7 +281,7 @@ def _get_environment_flags_from_document(self) -> Flags:
     def _get_identity_flags_from_document(
         self, identifier: str, traits: typing.Dict[str, typing.Any]
     ) -> Flags:
-        identity_model = self._build_identity_model(identifier, **traits)
+        identity_model = self._get_identity_model(identifier, **traits)
         feature_states = engine.get_identity_feature_states(
             self._environment, identity_model
         )
@@ -334,7 +346,11 @@ def _get_json_response(self, url: str, method: str, body: dict = None):
                 "Unable to get valid response from Flagsmith API."
             ) from e
 
-    def _build_identity_model(self, identifier: str, **traits):
+    def _get_identity_model(
+        self,
+        identifier: str,
+        **traits: TraitValue,
+    ) -> IdentityModel:
         if not self._environment:
             raise FlagsmithClientError(
                 "Unable to build identity model when no local environment present."
@@ -344,6 +360,11 @@ def _build_identity_model(self, identifier: str, **traits):
             TraitModel(trait_key=key, trait_value=value)
             for key, value in traits.items()
         ]
+
+        if identity := self._identity_overrides_by_identifier.get(identifier):
+            identity.update_traits(trait_models)
+            return identity
+
         return IdentityModel(
             identifier=identifier,
             environment_api_key=self._environment.api_key,

flagsmith/models.py

@@ -10,20 +10,19 @@
 @dataclass
 class BaseFlag:
     enabled: bool
-    value: typing.Union[str, int, float, bool, type(None)]
-    is_default: bool
+    value: typing.Union[str, int, float, bool, None]
 
 
+@dataclass
 class DefaultFlag(BaseFlag):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, is_default=True, **kwargs)
+    is_default: bool = field(default=True)
 
 
+@dataclass
 class Flag(BaseFlag):
-    def __init__(self, *args, feature_id: int, feature_name: str, **kwargs):
-        super().__init__(*args, is_default=False, **kwargs)
-        self.feature_id = feature_id
-        self.feature_name = feature_name
+    feature_id: int
+    feature_name: str
+    is_default: bool = field(default=False)
 
     @classmethod
     def from_feature_state_model(

poetry.lock

@@ -10,10 +10,10 @@ documentation = "https://docs.flagsmith.com"
 packages = [{ include = "flagsmith" }]
 
 [tool.poetry.dependencies]
-python = ">=3.7.0,<4"
+python = ">=3.8.0,<4"
 requests = "^2.27.1"
 requests-futures = "^1.0.0"
-flagsmith-flag-engine = "^5.0.0"
+flagsmith-flag-engine = "^5.1.0"
 sseclient-py = "^1.8.0"
 pytz = "^2023.4"
 
@@ -28,6 +28,7 @@ isort = "^5.10.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.0"
+pytest-cov = "^4.1.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

pyproject.toml

@@ -53,5 +53,30 @@
             "enabled": true
         }
     ],
-    "updated_at": "2023-07-14 16:12:00.000000"
+    "updated_at": "2023-07-14 16:12:00.000000",
+    "identity_overrides": [
+        {
+            "identifier": "overridden-id",
+            "identity_uuid": "0f21cde8-63c5-4e50-baca-87897fa6cd01",
+            "created_date": "2019-08-27T14:53:45.698555Z",
+            "updated_at": "2023-07-14 16:12:00.000000",
+            "environment_api_key": "B62qaMZNwfiqT76p38ggrQ",
+            "identity_features": [
+                {
+                    "id": 1,
+                    "feature": {
+                        "id": 1,
+                        "name": "some_feature",
+                        "type": "STANDARD"
+                    },
+                    "featurestate_uuid": "1bddb9a5-7e59-42c6-9be9-625fa369749f",
+                    "feature_state_value": "some-overridden-value",
+                    "enabled": false,
+                    "environment": 1,
+                    "identity": null,
+                    "feature_segment": null
+                }
+            ]
+        }
+    ]
 }

tests/data/environment.json

@@ -1,4 +1,5 @@
 import json
+import time
 import typing
 import uuid
 
@@ -512,3 +513,30 @@ def test_error_raised_when_realtime_updates_is_true_and_local_evaluation_false(
             enable_local_evaluation=False,
             enable_realtime_updates=True,
         )
+
+
+@responses.activate()
+def test_flagsmith_client_get_identity_flags__local_evaluation__returns_expected(
+    environment_json: str,
+    server_api_key: str,
+) -> None:
+    # Given
+    identifier = "overridden-id"
+
+    api_url = "https://mocked.flagsmith.com/api/v1/"
+    environment_document_url = f"{api_url}environment-document/"
+    responses.add(method="GET", url=environment_document_url, body=environment_json)
+
+    flagsmith = Flagsmith(
+        environment_key=server_api_key,
+        api_url=api_url,
+        enable_local_evaluation=True,
+    )
+    time.sleep(0.1)
+
+    # When
+    flag = flagsmith.get_identity_flags(identifier).get_flag("some_feature")
+
+    # Then
+    assert flag.enabled is False
+    assert flag.value == "some-overridden-value"