From 182ea04c7f776606c931f34c4fad32849f667d2c Mon Sep 17 00:00:00 2001 From: Matthew Elwell Date: Wed, 26 Jun 2024 14:54:46 +0100 Subject: [PATCH] feat(api-usage): add `subscription.plan` trait to `flagsmith.get_identity_flags` (#4247) --- api/organisations/tasks.py | 5 ++++- .../unit/organisations/test_unit_organisations_tasks.py | 6 +++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index 63ade5f5cfbc..ef60d76acf89 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -211,7 +211,10 @@ def handle_api_usage_notifications() -> None: ): feature_enabled = flagsmith_client.get_identity_flags( organisation.flagsmith_identifier, - traits={"organisation_id": organisation.id}, + traits={ + "organisation_id": organisation.id, + "subscription.plan": organisation.subscription.plan, + }, ).is_feature_enabled("api_usage_alerting") if not feature_enabled: continue diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 375b55b773c8..accd8c2327d2 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -272,7 +272,11 @@ def test_handle_api_usage_notifications_when_feature_flag_is_off( mock_api_usage.assert_not_called() client_mock.get_identity_flags.assert_called_once_with( - organisation.flagsmith_identifier, traits={"organisation_id": organisation.id} + organisation.flagsmith_identifier, + traits={ + "organisation_id": organisation.id, + "subscription.plan": organisation.subscription.plan, + }, ) assert len(mailoutbox) == 0