From 5ee6ad7d84ebd820790335f98e31434aa1bc709b Mon Sep 17 00:00:00 2001
From: Lendemor <thomas.brandeho@gmail.com>
Date: Fri, 14 Feb 2025 20:14:09 +0100
Subject: [PATCH] update integrations tests

---
 tests/integration/test_connection_banner.py | 37 ++++++++++++---------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/tests/integration/test_connection_banner.py b/tests/integration/test_connection_banner.py
index bfc9ea0ae..5be4c08fb 100644
--- a/tests/integration/test_connection_banner.py
+++ b/tests/integration/test_connection_banner.py
@@ -7,23 +7,27 @@ import pytest
 from selenium.common.exceptions import NoSuchElementException
 from selenium.webdriver.common.by import By
 
+from reflex import constants
 from reflex.testing import AppHarness, WebDriver
 
 from .utils import SessionStorage
 
 
-def ConnectionBanner(is_reflex_cloud: bool = False):
+def ConnectionBanner(
+    compile_context: constants.CompileContext = constants.CompileContext.RUN,
+):
     """App with a connection banner.
 
     Args:
-        is_reflex_cloud: The value for config.is_reflex_cloud.
+        compile_context: The value used to set the compile context.
     """
     import asyncio
 
     import reflex as rx
+    from reflex.config import environment
 
     # Simulate reflex cloud deploy
-    rx.config.get_config().is_reflex_cloud = is_reflex_cloud
+    environment.REFLEX_COMPILE_CONTEXT.set(compile_context)
 
     class State(rx.State):
         foo: int = 0
@@ -49,16 +53,17 @@ def ConnectionBanner(is_reflex_cloud: bool = False):
 
 
 @pytest.fixture(
-    params=[False, True], ids=["reflex_cloud_disabled", "reflex_cloud_enabled"]
+    params=[constants.CompileContext.RUN, constants.CompileContext.DEPLOY],
+    ids=["compile_context_run", "compile_context_deploy"],
 )
-def simulate_is_reflex_cloud(request) -> bool:
+def simulate_compile_context(request) -> constants.CompileContext:
     """Fixture to simulate reflex cloud deployment.
 
     Args:
         request: pytest request fixture.
 
     Returns:
-        True if reflex cloud is enabled, False otherwise.
+        The context to run the app with.
     """
     return request.param
 
@@ -66,13 +71,13 @@ def simulate_is_reflex_cloud(request) -> bool:
 @pytest.fixture()
 def connection_banner(
     tmp_path,
-    simulate_is_reflex_cloud: bool,
+    simulate_compile_context: constants.CompileContext,
 ) -> Generator[AppHarness, None, None]:
     """Start ConnectionBanner app at tmp_path via AppHarness.
 
     Args:
         tmp_path: pytest tmp_path fixture
-        simulate_is_reflex_cloud: Whether is_reflex_cloud is set for the app.
+        simulate_compile_context: Which context to run the app with.
 
     Yields:
         running AppHarness instance
@@ -80,11 +85,13 @@ def connection_banner(
     with AppHarness.create(
         root=tmp_path,
         app_source=functools.partial(
-            ConnectionBanner, is_reflex_cloud=simulate_is_reflex_cloud
+            ConnectionBanner, compile_context=simulate_compile_context
+        ),
+        app_name=(
+            "connection_banner_reflex_cloud"
+            if simulate_compile_context == constants.CompileContext.DEPLOY
+            else "connection_banner"
         ),
-        app_name="connection_banner_reflex_cloud"
-        if simulate_is_reflex_cloud
-        else "connection_banner",
     ) as harness:
         yield harness
 
@@ -194,13 +201,13 @@ async def test_connection_banner(connection_banner: AppHarness):
 
 @pytest.mark.asyncio
 async def test_cloud_banner(
-    connection_banner: AppHarness, simulate_is_reflex_cloud: bool
+    connection_banner: AppHarness, simulate_compile_context: constants.CompileContext
 ):
     """Test that the connection banner is displayed when the websocket drops.
 
     Args:
         connection_banner: AppHarness instance.
-        simulate_is_reflex_cloud: Whether is_reflex_cloud is set for the app.
+        simulate_compile_context: Which context to set for the app.
     """
     assert connection_banner.app_instance is not None
     assert connection_banner.backend is not None
@@ -213,7 +220,7 @@ async def test_cloud_banner(
 
     driver.add_cookie({"name": "backend-enabled", "value": "false"})
     driver.refresh()
-    if simulate_is_reflex_cloud:
+    if simulate_compile_context == constants.CompileContext.DEPLOY:
         assert connection_banner._poll_for(lambda: has_cloud_banner(driver))
     else:
         _assert_token(connection_banner, driver)