Fix Session context manager mocking in pre-start tests

Author: sajalshh

backend/tests/scripts/test_backend_pre_start.py

@@ -8,23 +8,16 @@
 def test_init_successful_connection() -> None:
     engine_mock = MagicMock()
 
-    session_mock = MagicMock()
-    session_mock.exec.return_value = True
+    session_cm = MagicMock()
+    session = session_cm.__enter__.return_value
+    session.exec.return_value = True
 
     with (
-        patch("sqlmodel.Session", return_value=session_mock),
+        patch("sqlmodel.Session", return_value=session_cm),
         patch.object(logger, "info"),
         patch.object(logger, "error"),
         patch.object(logger, "warn"),
     ):
-        try:
-            init(engine_mock)
-            connection_successful = True
-        except Exception:
-            connection_successful = False
+        init(engine_mock)
 
-        assert connection_successful, (
-            "The database connection should be successful and not raise an exception."
-        )
-
-        session_mock.exec.assert_called_once_with(select(1))
+        session.exec.assert_called_once_with(select(1))

backend/tests/scripts/test_test_pre_start.py

@@ -8,23 +8,16 @@
 def test_init_successful_connection() -> None:
     engine_mock = MagicMock()
 
-    session_mock = MagicMock()
-    session_mock.exec.return_value = True
+    session_cm = MagicMock()
+    session = session_cm.__enter__.return_value
+    session.exec.return_value = True
 
     with (
-        patch("sqlmodel.Session", return_value=session_mock),
+        patch("sqlmodel.Session", return_value=session_cm),
         patch.object(logger, "info"),
         patch.object(logger, "error"),
         patch.object(logger, "warn"),
     ):
-        try:
-            init(engine_mock)
-            connection_successful = True
-        except Exception:
-            connection_successful = False
+        init(engine_mock)
 
-        assert connection_successful, (
-            "The database connection should be successful and not raise an exception."
-        )
-
-        session_mock.exec.assert_called_once_with(select(1))
+        session.exec.assert_called_once_with(select(1))