From 09fa1c1a6188139e3e64f7e92b918ce351abff4d Mon Sep 17 00:00:00 2001
From: Lendemor <thomas.brandeho@gmail.com>
Date: Thu, 19 Dec 2024 23:39:53 +0100
Subject: [PATCH] fix tests

---
 reflex/app.py                       |  4 +--
 reflex/utils/prerequisites.py       | 18 ++++++++++
 tests/units/test_health_endpoint.py | 52 +++++++++++++++++++++--------
 3 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/reflex/app.py b/reflex/app.py
index a5a6dfe6b..032b198f6 100644
--- a/reflex/app.py
+++ b/reflex/app.py
@@ -1358,9 +1358,9 @@ async def health() -> JSONResponse:
 
     tasks = []
 
-    if get_config().db_url:
+    if prerequisites.check_db_used():
         tasks.append(get_db_status())
-    if get_config().redis_url:
+    if prerequisites.check_redis_used():
         tasks.append(prerequisites.get_redis_status())
 
     results = await asyncio.gather(*tasks)
diff --git a/reflex/utils/prerequisites.py b/reflex/utils/prerequisites.py
index 97c6a5974..ca40f8e6c 100644
--- a/reflex/utils/prerequisites.py
+++ b/reflex/utils/prerequisites.py
@@ -1174,6 +1174,24 @@ def initialize_frontend_dependencies():
     initialize_web_directory()
 
 
+def check_db_used() -> bool:
+    """Check if the database is used.
+
+    Returns:
+        True if the database is used.
+    """
+    return bool(get_config().db_url)
+
+
+def check_redis_used() -> bool:
+    """Check if Redis is used.
+
+    Returns:
+        True if Redis is used.
+    """
+    return bool(get_config().redis_url)
+
+
 def check_db_initialized() -> bool:
     """Check if the database migrations are initialized.
 
diff --git a/tests/units/test_health_endpoint.py b/tests/units/test_health_endpoint.py
index fe350266f..6d12d79d6 100644
--- a/tests/units/test_health_endpoint.py
+++ b/tests/units/test_health_endpoint.py
@@ -15,11 +15,11 @@ from reflex.utils.prerequisites import get_redis_status
     "mock_redis_client, expected_status",
     [
         # Case 1: Redis client is available and responds to ping
-        (Mock(ping=lambda: None), True),
+        (Mock(ping=lambda: None), {"redis": True}),
         # Case 2: Redis client raises RedisError
-        (Mock(ping=lambda: (_ for _ in ()).throw(RedisError)), False),
+        (Mock(ping=lambda: (_ for _ in ()).throw(RedisError)), {"redis": False}),
         # Case 3: Redis client is not used
-        (None, None),
+        (None, {"redis": None}),
     ],
 )
 async def test_get_redis_status(mock_redis_client, expected_status, mocker):
@@ -41,12 +41,12 @@ async def test_get_redis_status(mock_redis_client, expected_status, mocker):
     "mock_engine, execute_side_effect, expected_status",
     [
         # Case 1: Database is accessible
-        (MagicMock(), None, True),
+        (MagicMock(), None, {"db": True}),
         # Case 2: Database connection error (OperationalError)
         (
             MagicMock(),
             sqlalchemy.exc.OperationalError("error", "error", "error"),
-            False,
+            {"db": False},
         ),
     ],
 )
@@ -74,25 +74,49 @@ async def test_get_db_status(mock_engine, execute_side_effect, expected_status,
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    "db_status, redis_status, expected_status, expected_code",
+    "db_enabled, redis_enabled, db_status, redis_status, expected_status, expected_code",
     [
         # Case 1: Both services are connected
-        (True, True, {"status": True, "db": True, "redis": True}, 200),
+        (True, True, True, True, {"status": True, "db": True, "redis": True}, 200),
         # Case 2: Database not connected, Redis connected
-        (False, True, {"status": False, "db": False, "redis": True}, 503),
+        (True, True, False, True, {"status": False, "db": False, "redis": True}, 503),
         # Case 3: Database connected, Redis not connected
-        (True, False, {"status": False, "db": True, "redis": False}, 503),
+        (True, True, True, False, {"status": False, "db": True, "redis": False}, 503),
         # Case 4: Both services not connected
-        (False, False, {"status": False, "db": False, "redis": False}, 503),
+        (True, True, False, False, {"status": False, "db": False, "redis": False}, 503),
         # Case 5: Database Connected, Redis not used
-        (True, None, {"status": True, "db": True, "redis": False}, 200),
+        (True, False, True, None, {"status": True, "db": True}, 200),
+        # Case 6: Database not used, Redis Connected
+        (False, True, None, True, {"status": True, "redis": True}, 200),
+        # Case 7: Both services not used
+        (False, False, None, None, {"status": True}, 200),
     ],
 )
-async def test_health(db_status, redis_status, expected_status, expected_code, mocker):
+async def test_health(
+    db_enabled,
+    redis_enabled,
+    db_status,
+    redis_status,
+    expected_status,
+    expected_code,
+    mocker,
+):
     # Mock get_db_status and get_redis_status
-    mocker.patch("reflex.app.get_db_status", return_value=db_status)
     mocker.patch(
-        "reflex.utils.prerequisites.get_redis_status", return_value=redis_status
+        "reflex.utils.prerequisites.check_db_used",
+        return_value=db_enabled,
+    )
+    mocker.patch(
+        "reflex.utils.prerequisites.check_redis_used",
+        return_value=redis_enabled,
+    )
+    mocker.patch(
+        "reflex.app.get_db_status",
+        return_value={"db": db_status},
+    )
+    mocker.patch(
+        "reflex.utils.prerequisites.get_redis_status",
+        return_value={"redis": redis_status},
     )
 
     # Call the async health function