From 0be2c3e810125d00f631b848f10ba81f3cf6a900 Mon Sep 17 00:00:00 2001
From: Benedikt Bartscher <bb@senbax.de>
Date: Thu, 21 Nov 2024 02:13:41 +0100
Subject: [PATCH] fix: properly cleanup env vars with pytest fixtures

---
 reflex/testing.py             | 12 ------------
 tests/integration/conftest.py | 13 +++++++++----
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/reflex/testing.py b/reflex/testing.py
index 5223beda6..036cf219a 100644
--- a/reflex/testing.py
+++ b/reflex/testing.py
@@ -975,17 +975,6 @@ class AppHarnessProd(AppHarness):
         finally:
             EnvironmentVariables.REFLEX_SKIP_COMPILE.set(None)
 
-    @override
-    def start(self) -> AppHarnessProd:
-        """Start AppHarnessProd instance.
-
-        Returns:
-            self
-        """
-        EnvironmentVariables.REFLEX_ENV_MODE.set(reflex.constants.base.Env.PROD)
-        _ = super().start()
-        return self
-
     @override
     def stop(self):
         """Stop the frontend python webserver."""
@@ -994,4 +983,3 @@ class AppHarnessProd(AppHarness):
             self.frontend_server.shutdown()
         if self.frontend_thread is not None:
             self.frontend_thread.join()
-        EnvironmentVariables.REFLEX_ENV_MODE.set(None)
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index 337f4e174..aaae9faeb 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -3,12 +3,13 @@
 import os
 import re
 from pathlib import Path
-from typing import Type
+from typing import Generator, Type
 
 import pytest
 
 import reflex.constants
 from reflex.config import EnvironmentVariables
+from reflex.constants.base import Env
 from reflex.testing import AppHarness, AppHarnessProd
 
 DISPLAY = None
@@ -81,14 +82,18 @@ def pytest_exception_interact(node, call, report):
 )
 def app_harness_env(
     request: pytest.FixtureRequest,
-) -> Type[AppHarness]:
+) -> Generator[Type[AppHarness], None, None]:
     """Parametrize the AppHarness class to use for the test, either dev or prod.
 
     Args:
         request: The pytest fixture request object.
 
-    Returns:
+    Yields:
         The AppHarness class to use for the test.
     """
     harness: Type[AppHarness] = request.param
-    return harness
+    if issubclass(harness, AppHarnessProd):
+        EnvironmentVariables.REFLEX_ENV_MODE.set(Env.PROD)
+    yield harness
+    if isinstance(harness, AppHarnessProd):
+        EnvironmentVariables.REFLEX_ENV_MODE.set(None)