Implementation of offline mode (single client class) (#50)

matthewelwell · web-flow · commit 36fd7e90277d · 2023-07-31T09:49:53.000+01:00
* Initial implementation of offline mode (single client class)

* Use offline handler if API request fails

* Improve test coverage
diff --git a/flagsmith/flagsmith.py b/flagsmith/flagsmith.py
@@ -14,6 +14,7 @@
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.exceptions import FlagsmithAPIError, FlagsmithClientError
 from flagsmith.models import DefaultFlag, Flags, Segment
+from flagsmith.offline_handlers import BaseOfflineHandler
 from flagsmith.polling_manager import EnvironmentDataPollingManager
 from flagsmith.utils.identities import generate_identities_data
 
@@ -39,8 +40,8 @@ class Flagsmith:
 
     def __init__(
         self,
-        environment_key: str,
-        api_url: str = DEFAULT_API_URL,
+        environment_key: str = None,
+        api_url: str = None,
         custom_headers: typing.Dict[str, typing.Any] = None,
         request_timeout_seconds: int = None,
         enable_local_evaluation: bool = False,
@@ -49,9 +50,12 @@ def __init__(
         enable_analytics: bool = False,
         default_flag_handler: typing.Callable[[str], DefaultFlag] = None,
         proxies: typing.Dict[str, str] = None,
+        offline_mode: bool = False,
+        offline_handler: BaseOfflineHandler = None,
     ):
         """
-        :param environment_key: The environment key obtained from Flagsmith interface
+        :param environment_key: The environment key obtained from Flagsmith interface.
+            Required unless offline_mode is True.
         :param api_url: Override the URL of the Flagsmith API to communicate with
         :param custom_headers: Additional headers to add to requests made to the
             Flagsmith API
@@ -65,59 +69,83 @@ def __init__(
         :param enable_analytics: if enabled, sends additional requests to the Flagsmith
             API to power flag analytics charts
         :param default_flag_handler: callable which will be used in the case where
-            flags cannot be retrieved from the API or a non existent feature is
+            flags cannot be retrieved from the API or a non-existent feature is
             requested
         :param proxies: as per https://requests.readthedocs.io/en/latest/api/#requests.Session.proxies
+        :param offline_mode: sets the client into offline mode. Relies on offline_handler for
+            evaluating flags.
+        :param offline_handler: provide a handler for offline logic. Used to get environment
+            document from another source when in offline_mode. Works in place of
+            default_flag_handler if offline_mode is not set and using remote evaluation.
         """
-        self.session = requests.Session()
-        self.session.headers.update(
-            **{"X-Environment-Key": environment_key}, **(custom_headers or {})
-        )
-        self.session.proxies.update(proxies or {})
-        retries = retries or Retry(total=3, backoff_factor=0.1)
-
-        self.api_url = api_url if api_url.endswith("/") else f"{api_url}/"
-        self.request_timeout_seconds = request_timeout_seconds
-        self.session.mount(self.api_url, HTTPAdapter(max_retries=retries))
-
-        self.environment_flags_url = f"{self.api_url}flags/"
-        self.identities_url = f"{self.api_url}identities/"
-        self.environment_url = f"{self.api_url}environment-document/"
 
+        self.offline_mode = offline_mode
+        self.enable_local_evaluation = enable_local_evaluation
+        self.offline_handler = offline_handler
+        self.default_flag_handler = default_flag_handler
+        self._analytics_processor = None
         self._environment = None
-        if enable_local_evaluation:
-            if not environment_key.startswith("ser."):
-                raise ValueError(
-                    "In order to use local evaluation, please generate a server key "
-                    "in the environment settings page."
-                )
 
-            self.environment_data_polling_manager_thread = (
-                EnvironmentDataPollingManager(
-                    main=self,
-                    refresh_interval_seconds=environment_refresh_interval_seconds,
-                    daemon=True,  # noqa
-                )
+        # argument validation
+        if offline_mode and not offline_handler:
+            raise ValueError("offline_handler must be provided to use offline mode.")
+        elif default_flag_handler and offline_handler:
+            raise ValueError(
+                "Cannot use both default_flag_handler and offline_handler."
             )
-            self.environment_data_polling_manager_thread.start()
 
-        self._analytics_processor = (
-            AnalyticsProcessor(
-                environment_key, self.api_url, timeout=self.request_timeout_seconds
+        if self.offline_handler:
+            self._environment = self.offline_handler.get_environment()
+
+        if not self.offline_mode:
+            if not environment_key:
+                raise ValueError("environment_key is required.")
+
+            self.session = requests.Session()
+            self.session.headers.update(
+                **{"X-Environment-Key": environment_key}, **(custom_headers or {})
             )
-            if enable_analytics
-            else None
-        )
+            self.session.proxies.update(proxies or {})
+            retries = retries or Retry(total=3, backoff_factor=0.1)
+
+            api_url = api_url or DEFAULT_API_URL
+            self.api_url = api_url if api_url.endswith("/") else f"{api_url}/"
+
+            self.request_timeout_seconds = request_timeout_seconds
+            self.session.mount(self.api_url, HTTPAdapter(max_retries=retries))
+
+            self.environment_flags_url = f"{self.api_url}flags/"
+            self.identities_url = f"{self.api_url}identities/"
+            self.environment_url = f"{self.api_url}environment-document/"
+
+            if self.enable_local_evaluation:
+                if not environment_key.startswith("ser."):
+                    raise ValueError(
+                        "In order to use local evaluation, please generate a server key "
+                        "in the environment settings page."
+                    )
+
+                self.environment_data_polling_manager_thread = (
+                    EnvironmentDataPollingManager(
+                        main=self,
+                        refresh_interval_seconds=environment_refresh_interval_seconds,
+                        daemon=True,  # noqa
+                    )
+                )
+                self.environment_data_polling_manager_thread.start()
 
-        self.default_flag_handler = default_flag_handler
+            if enable_analytics:
+                self._analytics_processor = AnalyticsProcessor(
+                    environment_key, self.api_url, timeout=self.request_timeout_seconds
+                )
 
     def get_environment_flags(self) -> Flags:
         """
         Get all the default for flags for the current environment.
 
         :return: Flags object holding all the flags for the current environment.
         """
-        if self._environment:
+        if (self.offline_mode or self.enable_local_evaluation) and self._environment:
             return self._get_environment_flags_from_document()
         return self._get_environment_flags_from_api()
 
@@ -136,7 +164,7 @@ def get_identity_flags(
         :return: Flags object holding all the flags for the given identity.
         """
         traits = traits or {}
-        if self._environment:
+        if (self.offline_mode or self.enable_local_evaluation) and self._environment:
             return self._get_identity_flags_from_document(identifier, traits)
         return self._get_identity_flags_from_api(identifier, traits)
 
@@ -202,7 +230,9 @@ def _get_environment_flags_from_api(self) -> Flags:
                 default_flag_handler=self.default_flag_handler,
             )
         except FlagsmithAPIError:
-            if self.default_flag_handler:
+            if self.offline_handler:
+                return self._get_environment_flags_from_document()
+            elif self.default_flag_handler:
                 return Flags(default_flag_handler=self.default_flag_handler)
             raise
 
@@ -220,7 +250,9 @@ def _get_identity_flags_from_api(
                 default_flag_handler=self.default_flag_handler,
             )
         except FlagsmithAPIError:
-            if self.default_flag_handler:
+            if self.offline_handler:
+                return self._get_identity_flags_from_document(identifier, traits)
+            elif self.default_flag_handler:
                 return Flags(default_flag_handler=self.default_flag_handler)
             raise
 
diff --git a/flagsmith/offline_handlers.py b/flagsmith/offline_handlers.py
@@ -0,0 +1,22 @@
+import json
+from abc import ABC, abstractmethod
+
+from flag_engine.environments.builders import build_environment_model
+from flag_engine.environments.models import EnvironmentModel
+
+
+class BaseOfflineHandler(ABC):
+    @abstractmethod
+    def get_environment(self) -> EnvironmentModel:
+        raise NotImplementedError()
+
+
+class LocalFileHandler(BaseOfflineHandler):
+    def __init__(self, environment_document_path: str):
+        with open(environment_document_path) as environment_document:
+            self.environment = build_environment_model(
+                json.loads(environment_document.read())
+            )
+
+    def get_environment(self) -> EnvironmentModel:
+        return self.environment
diff --git a/tests/test_flagsmith.py b/tests/test_flagsmith.py
@@ -1,4 +1,5 @@
 import json
+import typing
 import uuid
 
 import pytest
@@ -8,7 +9,12 @@
 
 from flagsmith import Flagsmith
 from flagsmith.exceptions import FlagsmithAPIError
-from flagsmith.models import DefaultFlag
+from flagsmith.models import DefaultFlag, Flags
+from flagsmith.offline_handlers import BaseOfflineHandler
+
+if typing.TYPE_CHECKING:
+    from flag_engine.environments.models import EnvironmentModel
+    from pytest_mock import MockerFixture
 
 
 def test_flagsmith_starts_polling_manager_on_init_if_enabled(mocker, server_api_key):
@@ -68,6 +74,7 @@ def test_get_environment_flags_uses_local_environment_when_available(
 ):
     # Given
     flagsmith._environment = environment_model
+    flagsmith.enable_local_evaluation = True
 
     # When
     all_flags = flagsmith.get_environment_flags().all_flags()
@@ -134,6 +141,7 @@ def test_get_identity_flags_uses_local_environment_when_available(
 ):
     # Given
     flagsmith._environment = environment_model
+    flagsmith.enable_local_evaluation = True
     mock_engine = mocker.patch("flagsmith.flagsmith.engine")
 
     feature_state = FeatureStateModel(
@@ -378,3 +386,91 @@ def test_initialise_flagsmith_with_proxies():
 
     # Then
     assert flagsmith.session.proxies == proxies
+
+
+def test_offline_mode(environment_model: "EnvironmentModel") -> None:
+    # Given
+    class DummyOfflineHandler(BaseOfflineHandler):
+        def get_environment(self) -> "EnvironmentModel":
+            return environment_model
+
+    # When
+    flagsmith = Flagsmith(offline_mode=True, offline_handler=DummyOfflineHandler())
+
+    # Then
+    # we can request the flags from the client successfully
+    environment_flags: Flags = flagsmith.get_environment_flags()
+    assert environment_flags.is_feature_enabled("some_feature") is True
+
+    identity_flags: Flags = flagsmith.get_identity_flags("identity")
+    assert identity_flags.is_feature_enabled("some_feature") is True
+
+
+@responses.activate()
+def test_flagsmith_uses_offline_handler_if_set_and_no_api_response(
+    mocker: "MockerFixture", environment_model: "EnvironmentModel"
+) -> None:
+    # Given
+    api_url = "http://some.flagsmith.com/api/v1/"
+    mock_offline_handler = mocker.MagicMock(spec=BaseOfflineHandler)
+    mock_offline_handler.get_environment.return_value = environment_model
+
+    flagsmith = Flagsmith(
+        environment_key="some-key",
+        api_url=api_url,
+        offline_handler=mock_offline_handler,
+    )
+
+    responses.add(flagsmith.environment_flags_url, status=500)
+    responses.add(flagsmith.identities_url, status=500)
+
+    # When
+    environment_flags = flagsmith.get_environment_flags()
+    identity_flags = flagsmith.get_identity_flags("identity", traits={})
+
+    # Then
+    mock_offline_handler.get_environment.assert_called_once_with()
+
+    assert environment_flags.is_feature_enabled("some_feature") is True
+    assert environment_flags.get_feature_value("some_feature") == "some-value"
+
+    assert identity_flags.is_feature_enabled("some_feature") is True
+    assert identity_flags.get_feature_value("some_feature") == "some-value"
+
+
+def test_cannot_use_offline_mode_without_offline_handler():
+    with pytest.raises(ValueError) as e:
+        # When
+        Flagsmith(offline_mode=True, offline_handler=None)
+
+    # Then
+    assert (
+        e.exconly()
+        == "ValueError: offline_handler must be provided to use offline mode."
+    )
+
+
+def test_cannot_use_default_handler_and_offline_handler(mocker):
+    # When
+    with pytest.raises(ValueError) as e:
+        Flagsmith(
+            offline_handler=mocker.MagicMock(spec=BaseOfflineHandler),
+            default_flag_handler=lambda flag_name: DefaultFlag(
+                enabled=True, value="foo"
+            ),
+        )
+
+    # Then
+    assert (
+        e.exconly()
+        == "ValueError: Cannot use both default_flag_handler and offline_handler."
+    )
+
+
+def test_cannot_create_flagsmith_client_in_remote_evaluation_without_api_key():
+    # When
+    with pytest.raises(ValueError) as e:
+        Flagsmith()
+
+    # Then
+    assert e.exconly() == "ValueError: environment_key is required."
diff --git a/tests/test_offline_handlers.py b/tests/test_offline_handlers.py
@@ -0,0 +1,22 @@
+from unittest.mock import mock_open, patch
+
+from flag_engine.environments.models import EnvironmentModel
+
+from flagsmith.offline_handlers import LocalFileHandler
+
+
+def test_local_file_handler(environment_json):
+    with patch("builtins.open", mock_open(read_data=environment_json)) as mock_file:
+        # Given
+        environment_document_file_path = "/some/path/environment.json"
+        local_file_handler = LocalFileHandler(environment_document_file_path)
+
+        # When
+        environment_model = local_file_handler.get_environment()
+
+        # Then
+        assert isinstance(environment_model, EnvironmentModel)
+        assert (
+            environment_model.api_key == "B62qaMZNwfiqT76p38ggrQ"
+        )  # hard coded from json file
+        mock_file.assert_called_once_with(environment_document_file_path)
