From fd2094131e5e4972a17be5a8efa4cb379cc821ee Mon Sep 17 00:00:00 2001
From: Tushar <30565750+tushar5526@users.noreply.github.com>
Date: Tue, 16 Jul 2024 14:50:05 +0530
Subject: [PATCH] fix: Add a custom exception for invalid features (#86)

* fix: Add a custom exception for invalid features

* feat: add mock success response
---
 flagsmith/exceptions.py |  4 ++++
 flagsmith/models.py     |  6 ++++--
 tests/test_flagsmith.py | 26 ++++++++++++++++++++++----
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/flagsmith/exceptions.py b/flagsmith/exceptions.py
index c1c3b36..e499207 100644
--- a/flagsmith/exceptions.py
+++ b/flagsmith/exceptions.py
@@ -4,3 +4,7 @@ class FlagsmithClientError(Exception):
 
 class FlagsmithAPIError(FlagsmithClientError):
     pass
+
+
+class FlagsmithFeatureDoesNotExistError(FlagsmithClientError):
+    pass
diff --git a/flagsmith/models.py b/flagsmith/models.py
index d06a95f..ac7d2cc 100644
--- a/flagsmith/models.py
+++ b/flagsmith/models.py
@@ -6,7 +6,7 @@
 from flag_engine.features.models import FeatureStateModel
 
 from flagsmith.analytics import AnalyticsProcessor
-from flagsmith.exceptions import FlagsmithClientError
+from flagsmith.exceptions import FlagsmithFeatureDoesNotExistError
 
 
 @dataclass
@@ -135,7 +135,9 @@ def get_flag(self, feature_name: str) -> typing.Union[DefaultFlag, Flag]:
         except KeyError:
             if self.default_flag_handler:
                 return self.default_flag_handler(feature_name)
-            raise FlagsmithClientError("Feature does not exist: %s" % feature_name)
+            raise FlagsmithFeatureDoesNotExistError(
+                "Feature does not exist: %s" % feature_name
+            )
 
         if self._analytics_processor and hasattr(flag, "feature_name"):
             self._analytics_processor.track_feature(flag.feature_name)
diff --git a/tests/test_flagsmith.py b/tests/test_flagsmith.py
index 3ecc0fc..093aa88 100644
--- a/tests/test_flagsmith.py
+++ b/tests/test_flagsmith.py
@@ -11,15 +11,16 @@
 from pytest_mock import MockerFixture
 
 from flagsmith import Flagsmith
-from flagsmith.exceptions import FlagsmithAPIError
+from flagsmith.exceptions import (
+    FlagsmithAPIError,
+    FlagsmithFeatureDoesNotExistError,
+)
 from flagsmith.models import DefaultFlag, Flags
 from flagsmith.offline_handlers import BaseOfflineHandler
 
 
 def test_flagsmith_starts_polling_manager_on_init_if_enabled(
-    mocker: MockerFixture,
-    server_api_key: str,
-    requests_session_response_ok: None,
+    mocker: MockerFixture, server_api_key: str, requests_session_response_ok: None
 ) -> None:
     # Given
     mock_polling_manager = mocker.MagicMock()
@@ -559,3 +560,20 @@ def test_flagsmith_client_get_identity_flags__local_evaluation__returns_expected
     # Then
     assert flag.enabled is False
     assert flag.value == "some-overridden-value"
+
+
+def test_custom_feature_error_raised_when_invalid_feature(
+    requests_session_response_ok: None, server_api_key: str
+) -> None:
+    # Given
+    flagsmith = Flagsmith(
+        environment_key=server_api_key,
+        enable_local_evaluation=True,
+    )
+
+    flags = flagsmith.get_environment_flags()
+
+    # Then
+    with pytest.raises(FlagsmithFeatureDoesNotExistError):
+        # When
+        flags.is_feature_enabled("non-existing-feature")