Merge pull request #1838 from Agenta-AI/resolve-failing-backend-tests

[Enhancement]: Resolve Failing Backend Tests
Agenta-AI · Jul 3, 2024 · 2e3ee66 · 2e3ee66
2 parents 03ff817 + a8c0551
commit 2e3ee66
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 16 deletions.
diff --git a/agenta-backend/agenta_backend/models/db_engine.py b/agenta-backend/agenta_backend/models/db_engine.py
@@ -95,8 +95,15 @@ class DBEngine:
 
     def __init__(self) -> None:
         self.mode = os.environ.get("DATABASE_MODE", "v2")
-        self.postgres_uri = os.environ.get("POSTGRES_URI", None)
+        self.postgres_uri = os.environ.get("POSTGRES_URI")
         self.mongo_uri = os.environ.get("MONGODB_URI")
+        self.engine = create_async_engine(url=self.postgres_uri)  # type: ignore
+        self.async_session_maker = async_sessionmaker(
+            bind=self.engine, class_=AsyncSession, expire_on_commit=False
+        )
+        self.async_session = async_scoped_session(
+            session_factory=self.async_session_maker, scopefunc=current_task
+        )
 
     async def initialize_async_postgres(self):
         """
@@ -106,14 +113,6 @@ async def initialize_async_postgres(self):
         if not self.postgres_uri:
             raise ValueError("Postgres URI cannot be None.")
 
-        self.engine = create_async_engine(self.postgres_uri)
-        self.async_session_maker = async_sessionmaker(
-            bind=self.engine, class_=AsyncSession, expire_on_commit=False
-        )
-        self.async_session = async_scoped_session(
-            session_factory=self.async_session_maker, scopefunc=current_task
-        )
-
         async with self.engine.begin() as conn:
             # Drop and create tables if needed
             for model in models:
@@ -184,9 +183,5 @@ async def close(self):
 
         await self.engine.dispose()
 
-        self.engine = None
-        self.async_session_maker = None
-        self.async_session = None
-
 
 db_engine = DBEngine()
diff --git a/agenta-backend/agenta_backend/tests/conftest.py b/agenta-backend/agenta_backend/tests/conftest.py
@@ -16,10 +16,11 @@ def event_loop():
     res._close = res.close  # type: ignore
 
     # Initialize database and create tables
-    res.run_until_complete(DBEngine().init_db())
+    db_engine = DBEngine()
+    res.run_until_complete(db_engine.init_db())
 
     yield res
 
-    res.run_until_complete(DBEngine().close())  # close connections to database
-    res.run_until_complete(DBEngine().remove_db())  # drop database
+    res.run_until_complete(db_engine.remove_db())  # drop database
+    res.run_until_complete(db_engine.close())  # close connections to database
     res._close()  # close event loop # type: ignore