From 5fc5559858d35730275206af1c8c7ef60aa072b4 Mon Sep 17 00:00:00 2001
From: Zach Aysan <zachaysan@gmail.com>
Date: Thu, 14 Mar 2024 05:01:15 -0400
Subject: [PATCH] fix: Set the environment for local evaluation mode on init
 (#76)

* Set the environment for local evaluation mode on init

* Make tests pass with the update environment change

* Fix test failure

* Make update environment update local document

* Use responses for polling manager tests

* Lint

---------

Co-authored-by: Matthew Elwell <matthew.elwell@flagsmith.com>
---
 flagsmith/flagsmith.py        |  3 +++
 tests/test_flagsmith.py       |  4 +++-
 tests/test_polling_manager.py | 28 ++++++++++++++++++++++------
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/flagsmith/flagsmith.py b/flagsmith/flagsmith.py
index 3cd4c2b..bbfd227 100644
--- a/flagsmith/flagsmith.py
+++ b/flagsmith/flagsmith.py
@@ -185,6 +185,9 @@ def _initialise_local_evaluation(self) -> None:
             self.event_stream_thread.start()
 
         else:
+            # To ensure that the environment is set before allowing subsequent
+            # method calls, update the environment manually.
+            self.update_environment()
             self.environment_data_polling_manager_thread = (
                 EnvironmentDataPollingManager(
                     main=self,
diff --git a/tests/test_flagsmith.py b/tests/test_flagsmith.py
index f6afdbc..3ecc0fc 100644
--- a/tests/test_flagsmith.py
+++ b/tests/test_flagsmith.py
@@ -17,7 +17,9 @@
 
 
 def test_flagsmith_starts_polling_manager_on_init_if_enabled(
-    mocker: MockerFixture, server_api_key: str
+    mocker: MockerFixture,
+    server_api_key: str,
+    requests_session_response_ok: None,
 ) -> None:
     # Given
     mock_polling_manager = mocker.MagicMock()
diff --git a/tests/test_polling_manager.py b/tests/test_polling_manager.py
index c30df6f..e5ccb2d 100644
--- a/tests/test_polling_manager.py
+++ b/tests/test_polling_manager.py
@@ -2,6 +2,7 @@
 from unittest import mock
 
 import requests
+import responses
 from pytest_mock import MockerFixture
 
 from flagsmith import Flagsmith
@@ -41,17 +42,22 @@ def test_polling_manager_calls_update_environment_on_each_refresh() -> None:
     polling_manager.stop()
 
 
+@responses.activate()
 def test_polling_manager_is_resilient_to_api_errors(
-    mocker: MockerFixture, server_api_key: str
+    flagsmith: Flagsmith,
+    environment_json: str,
+    mocker: MockerFixture,
+    server_api_key: str,
 ) -> None:
     # Given
-    session_mock = mocker.patch("requests.Session")
-    session_mock.get.return_value = mock.MagicMock(status_code=500)
+    responses.add(method="GET", url=flagsmith.environment_url, body=environment_json)
     flagsmith = Flagsmith(
         environment_key=server_api_key,
         enable_local_evaluation=True,
         environment_refresh_interval_seconds=0.1,
     )
+
+    responses.add(method="GET", url=flagsmith.environment_url, status=500)
     polling_manager = flagsmith.environment_data_polling_manager_thread
 
     # Then
@@ -59,17 +65,27 @@ def test_polling_manager_is_resilient_to_api_errors(
     polling_manager.stop()
 
 
+@responses.activate()
 def test_polling_manager_is_resilient_to_request_exceptions(
-    mocker: MockerFixture, server_api_key: str
+    flagsmith: Flagsmith,
+    environment_json: str,
+    mocker: MockerFixture,
+    server_api_key: str,
 ) -> None:
     # Given
-    session_mock = mocker.patch("requests.Session")
-    session_mock.get.side_effect = requests.RequestException()
+    responses.add(method="GET", url=flagsmith.environment_url, body=environment_json)
     flagsmith = Flagsmith(
         environment_key=server_api_key,
         enable_local_evaluation=True,
         environment_refresh_interval_seconds=0.1,
     )
+
+    responses.add(
+        method="GET",
+        url=flagsmith.environment_url,
+        body=requests.RequestException("Some exception"),
+        status=500,
+    )
     polling_manager = flagsmith.environment_data_polling_manager_thread
 
     # Then