Skip to content

feat: Support registering global traits at client initialization #268

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
gagantrivedi wants to merge 1 commit into from
Closed

feat: Support registering global traits at client initialization #268

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 26 additions & 5 deletions flagsmith/flagsmith.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ def __init__(
offline_handler: typing.Optional[OfflineHandler] = None,
enable_realtime_updates: bool = False,
application_metadata: typing.Optional[ApplicationMetadata] = None,
default_traits: typing.Optional[TraitMapping] = None,
):
"""
:param environment_key: The environment key obtained from Flagsmith interface.
Expand Down Expand Up @@ -117,6 +118,12 @@ def __init__(
default_flag_handler if offline_mode is not set and using remote evaluation.
:param enable_realtime_updates: Use real-time functionality via SSE as opposed to polling the API
:param application_metadata: Optional metadata about the client application.
:param default_traits: traits that apply to every identity evaluation and
event for the lifetime of this client, e.g.
`{"app_name": "my-app", "branch": "main"}`. Traits passed to an
individual call take precedence over a default trait of the same key.
Envelope traits you don't want persisted in a dictionary with
`"transient"` and `"value"` keys, as for per-call traits.
"""

self.offline_mode = offline_mode
Expand All @@ -125,6 +132,7 @@ def __init__(
self.offline_handler = offline_handler
self.default_flag_handler = default_flag_handler
self.enable_realtime_updates = enable_realtime_updates
self.default_traits: TraitMapping = default_traits or {}
self._analytics_processor: typing.Optional[AnalyticsProcessor] = None
self._event_processor: typing.Optional[EventProcessor] = None
self._local_evaluation: typing.Optional[LocalEvaluation] = None
Expand Down Expand Up @@ -207,6 +215,15 @@ def __init__(
def _ensure_trailing_slash(url: str) -> str:
return url if url.endswith("/") else f"{url}/"

def _merge_traits(self, traits: typing.Optional[TraitMapping]) -> TraitMapping:
"""Combine the client's default traits with traits given for a single
call. Per-call traits win on key conflicts."""
if not self.default_traits:
return traits or {}
if not traits:
return self.default_traits
return {**self.default_traits, **traits}

def _initialise_analytics(
self,
environment_key: str,
Expand Down Expand Up @@ -305,10 +322,11 @@ def get_identity_flags(
Flagsmith, e.g. `{"num_orders": 10}`. Envelope traits you don't want persisted
in a dictionary with `"transient"` and `"value"` keys, e.g.
`{"num_orders": 10, "color": {"value": "pink", "transient": True}}`.
Merged over the client's `default_traits`, if any.
:param transient: if `True`, the identity won't get persisted
:return: Flags object holding all the flags for the given identity.
"""
traits = traits or {}
traits = self._merge_traits(traits)
if (
self.offline_mode or self.enable_local_evaluation
) and self._local_evaluation:
Expand All @@ -330,7 +348,8 @@ def get_identity_segments(
: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}
Flagsmith, e.g. {"num_orders": 10}. Merged over the client's
`default_traits`, if any.
:return: list of Segment objects that the identity is part of.
"""
if not (local_evaluation := self._local_evaluation):
Expand All @@ -341,7 +360,7 @@ def get_identity_segments(
context = map_context_and_identity_data_to_context(
context=local_evaluation.context,
identifier=identifier,
traits=traits,
traits=self._merge_traits(traits),
)

evaluation_result = engine.get_evaluation_result(
Expand All @@ -366,6 +385,8 @@ def get_experiment_flag(
other outcome is logged and skipped to keep experimentation data
clean. A `DefaultFlag` served via the `default_flag_handler` counts
as the feature not existing.

`traits` are merged over the client's `default_traits`, if any.
"""
if not self._event_processor:
raise ValueError("Events must be enabled to use experiment flags.")
Expand Down Expand Up @@ -423,7 +444,7 @@ def track_event(
event=event,
identifier=identifier,
value=value,
traits=resolve_trait_values(traits),
traits=resolve_trait_values(self._merge_traits(traits)),
metadata=metadata,
)

Expand All @@ -449,7 +470,7 @@ def track_exposure_event(
feature_name=feature_name,
identifier=identifier,
value=value,
traits=resolve_trait_values(traits),
traits=resolve_trait_values(self._merge_traits(traits)),
metadata=metadata,
)

Expand Down
242 changes: 242 additions & 0 deletions tests/test_flagsmith.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1455,3 +1455,245 @@ def test_get_experiment_flag__disabled_feature__skips_exposure(
"Not sending $flag_exposure for feature checkout_v2: feature is disabled."
in caplog.messages
)


def test_default_traits__not_set__empty(api_key: str) -> None:
# Given / When
flagsmith = Flagsmith(environment_key=api_key)

# Then
assert flagsmith.default_traits == {}


@responses.activate()
def test_default_traits__get_identity_flags__sent_to_api(
api_key: str, identities_json: str
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=api_key,
default_traits={
"app_name": "my-app",
"branch": {"value": "main", "transient": True},
},
)
responses.add(
method="POST",
url=flagsmith.identities_url,
body=identities_json,
match=[
matchers.json_params_matcher(
{
"identifier": "identifier",
"traits": [
{"trait_key": "app_name", "trait_value": "my-app"},
{
"trait_key": "branch",
"trait_value": "main",
"transient": True,
},
{"trait_key": "num_orders", "trait_value": 10},
],
}
)
],
)

# When & Then
flagsmith.get_identity_flags("identifier", traits={"num_orders": 10})


@responses.activate()
def test_default_traits__get_identity_flags__per_call_traits_take_precedence(
api_key: str, identities_json: str
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=api_key,
default_traits={"app_name": "my-app"},
)
responses.add(
method="POST",
url=flagsmith.identities_url,
body=identities_json,
match=[
matchers.json_params_matcher(
{
"identifier": "identifier",
"traits": [
{"trait_key": "app_name", "trait_value": "another-app"},
],
}
)
],
)

# When & Then
flagsmith.get_identity_flags("identifier", traits={"app_name": "another-app"})


def test_default_traits__local_evaluation__included_in_context(
requests_session_response_ok: None, server_api_key: str, mocker: MockerFixture
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=server_api_key,
enable_local_evaluation=True,
environment_refresh_interval_seconds=0.1,
default_traits={
"app_name": "my-app",
"branch": {"value": "main", "transient": True},
},
)
spy = mocker.spy(engine, "get_evaluation_result")

# When
flagsmith.get_identity_flags("identifier", traits={"num_orders": 10}).all_flags()

# Then
context = spy.call_args.kwargs.get("context") or spy.call_args.args[0]
assert context["identity"]["traits"] == {
"app_name": "my-app",
"branch": "main",
"num_orders": 10,
}


def test_default_traits__get_identity_segments__matched_without_per_call_traits(
requests_session_response_ok: None, server_api_key: str
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=server_api_key,
enable_local_evaluation=True,
environment_refresh_interval_seconds=0.1,
default_traits={"foo": "bar"}, # obtained from data/environment.json
)

# When
segments = flagsmith.get_identity_segments("identifier")

# Then
assert [segment.name for segment in segments] == ["Test segment"]


def test_default_traits__get_identity_segments__per_call_traits_take_precedence(
requests_session_response_ok: None, server_api_key: str
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=server_api_key,
enable_local_evaluation=True,
environment_refresh_interval_seconds=0.1,
default_traits={"foo": "not-bar"},
)

# When
segments = flagsmith.get_identity_segments("identifier", traits={"foo": "bar"})

# Then
assert [segment.name for segment in segments] == ["Test segment"]


def test_default_traits__track_event__merged_with_call_traits(
mocker: MockerFixture, api_key: str
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=api_key,
enable_events=True,
event_processor_config=EventProcessorConfig(events_api_url="http://test/"),
default_traits={
"app_name": "my-app",
"plan": {"value": "free", "transient": True},
},
)
mock_track = mocker.patch.object(flagsmith._event_processor, "track_event")

# When
flagsmith.track_event("purchase", identifier="user1", traits={"plan": "premium"})

# Then
mock_track.assert_called_once_with(
event="purchase",
identifier="user1",
value=None,
traits={"app_name": "my-app", "plan": "premium"},
metadata=None,
)


def test_default_traits__track_exposure_event__merged_with_call_traits(
mocker: MockerFixture, api_key: str
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=api_key,
enable_events=True,
event_processor_config=EventProcessorConfig(events_api_url="http://test/"),
default_traits={"app_name": "my-app"},
)
mock_track = mocker.patch.object(flagsmith._event_processor, "track_exposure_event")

# When
flagsmith.track_exposure_event(
"checkout_v2",
identifier="user1",
value="variant_b",
traits={"plan": "premium"},
)

# Then
mock_track.assert_called_once_with(
feature_name="checkout_v2",
identifier="user1",
value="variant_b",
traits={"app_name": "my-app", "plan": "premium"},
metadata=None,
)


@responses.activate()
def test_default_traits__get_experiment_flag__merged_once_per_call(
mocker: MockerFixture, api_key: str, identities_json: str
) -> None:
# Given
flagsmith = Flagsmith(
environment_key=api_key,
enable_events=True,
event_processor_config=EventProcessorConfig(events_api_url="http://test/"),
default_traits={"app_name": "my-app"},
)
mock_track = mocker.patch.object(flagsmith._event_processor, "track_exposure_event")
responses.add(
method="POST",
url=flagsmith.identities_url,
body=identities_json,
match=[
matchers.json_params_matcher(
{
"identifier": "user1",
"traits": [
{"trait_key": "app_name", "trait_value": "my-app"},
{"trait_key": "plan", "trait_value": "premium"},
],
}
)
],
)

# When
flagsmith.get_experiment_flag(
feature_name="some_feature",
identifier="user1",
traits={"plan": "premium"},
)

# Then
mock_track.assert_called_once_with(
feature_name="some_feature",
identifier="user1",
value="treatment",
traits={"app_name": "my-app", "plan": "premium"},
metadata={"experiment_id": 42},
)
Loading