TLK-2622 Unit tests improvements (#907)

.github/workflows/backend_unit_tests.yml

@@ -39,7 +39,7 @@ jobs:
       - name: Test with pytest
         if: github.actor != 'dependabot[bot]'
         run: |
-          make run-unit-tests-debug
+          make run-unit-tests
         env:
           PYTHONPATH: src
       - name: Upload coverage reports to Codecov

src/backend/alembic/env.py

@@ -16,7 +16,8 @@
 config = context.config
 
 # Overwrite alembic.file `sqlachemy.url` value
-config.set_main_option("sqlalchemy.url", Settings().get('database.url'))
+if not config.get_main_option("sqlalchemy.url"):
+    config.set_main_option("sqlalchemy.url", Settings().get('database.url'))
 
 # Interpret the config file for Python logging.
 # This line sets up loggers basically.

src/backend/main.py

@@ -57,7 +57,7 @@ async def lifespan(app: FastAPI):
     # Shutdown logic
 
 
-def create_app():
+def create_app() -> FastAPI:
     app = FastAPI(
         title="Cohere Toolkit API",
         description="This is the API for the Open Source Cohere Toolkit",

src/backend/pycharm_debug_main.py

@@ -50,7 +50,7 @@ async def lifespan(app: FastAPI):
     # Shutdown logic
 
 
-def create_app():
+def create_app() -> FastAPI:
     app = FastAPI(lifespan=lifespan)
 
     routers = [

src/backend/pytest.ini

@@ -1,5 +1,6 @@
 [pytest]
 env =
-    DATABASE_URL=postgresql://postgres:postgres@localhost:5433/postgres
+    DATABASE_URL=postgresql://postgres:postgres@localhost:5433
 filterwarnings =
     ignore::UserWarning:pydantic.*
+    ignore::DeprecationWarning

src/backend/tests/integration/routers/test_chat.py

@@ -24,7 +24,7 @@
 )
 
 
-@pytest.fixture()
+@pytest.fixture(scope="session")
 def user(session_chat: Session) -> User:
     return get_factory("User", session_chat).create()
 

src/backend/tests/unit/conftest.py

@@ -6,10 +6,12 @@
 import pytest
 from alembic.command import upgrade
 from alembic.config import Config
+from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from redis import Redis
 from sqlalchemy import create_engine
 from sqlalchemy.orm import Session
+from sqlalchemy.sql import text
 
 from backend.database_models import get_session
 from backend.database_models.base import CustomFilterQuery
@@ -19,23 +21,64 @@
 from backend.tests.unit.factories import get_factory
 
 DATABASE_URL = os.environ["DATABASE_URL"]
+MASTER_DB_NAME = "postgres"
+TEST_DB_PREFIX = "postgres_"
+MASTER_DATABASE_FULL_URL = f"{DATABASE_URL}/{MASTER_DB_NAME}"
+
+
+def create_test_database(test_db_name: str):
+    engine = create_engine(
+        MASTER_DATABASE_FULL_URL, echo=True, isolation_level="AUTOCOMMIT"
+    )
+
+    with engine.connect() as connection:
+        connection.execute(text(f"CREATE DATABASE {test_db_name}"))
+    engine.dispose()
+
+
+def drop_test_database_if_exists(test_db_name: str):
+    engine = create_engine(
+        MASTER_DATABASE_FULL_URL, echo=True, isolation_level="AUTOCOMMIT"
+    )
+
+    with engine.connect() as connection:
+        connection.execute(text(f"DROP DATABASE IF EXISTS {test_db_name}"))
+    engine.dispose()
+
+
+@pytest.fixture(scope="session")
+def fastapi_app() -> Generator[FastAPI, None, None]:
+    """Creates a session-scoped FastAPI app object."""
+    app = create_app()
+    yield app
 
 
 @pytest.fixture
 def client():
     yield TestClient(app)
 
 
-@pytest.fixture(scope="function")
-def engine() -> Generator[Any, None, None]:
+@pytest.fixture(scope="session")
+def engine(worker_id: str) -> Generator[Any, None, None]:
     """
     Yields a SQLAlchemy engine which is disposed of after the test session
     """
-    engine = create_engine(DATABASE_URL, echo=True)
+    test_db_name = f"{TEST_DB_PREFIX}{worker_id.replace('gw', '')}"
+    test_db_url = f"{DATABASE_URL}/{test_db_name}"
+
+    drop_test_database_if_exists(test_db_name)
+    create_test_database(test_db_name)
+    engine = create_engine(test_db_url, echo=True)
+
+    with engine.begin():
+        alembic_cfg = Config("src/backend/alembic.ini")
+        alembic_cfg.set_main_option("sqlalchemy.url", test_db_url)
+        upgrade(alembic_cfg, "head")
 
     yield engine
 
     engine.dispose()
+    drop_test_database_if_exists(test_db_name)
 
 
 @pytest.fixture(scope="function")
@@ -45,13 +88,10 @@ def session(engine: Any) -> Generator[Session, None, None]:
     that is rolled back after every function
     """
     connection = engine.connect()
-    # Begin the nested transaction
+    # Begin the transaction
     transaction = connection.begin()
     # Use connection within the started transaction
     session = Session(bind=connection, query_cls=CustomFilterQuery)
-    # Run Alembic migrations
-    alembic_cfg = Config("src/backend/alembic.ini")
-    upgrade(alembic_cfg, "head")
 
     yield session
 
@@ -63,55 +103,37 @@ def session(engine: Any) -> Generator[Session, None, None]:
 
 
 @pytest.fixture(scope="function")
-def session_client(session: Session) -> Generator[TestClient, None, None]:
+def session_client(session: Session, fastapi_app: FastAPI) -> Generator[TestClient, None, None]:
     """
     Fixture to inject the session into the API client
     """
 
     def override_get_session() -> Generator[Session, Any, None]:
         yield session
 
-    app = create_app()
-
-    app.dependency_overrides[get_session] = override_get_session
+    fastapi_app.dependency_overrides[get_session] = override_get_session
 
     print("Session at fixture " + str(session))
 
-    with TestClient(app) as client:
+    with TestClient(fastapi_app) as client:
         yield client
 
-    app.dependency_overrides = {}
+    fastapi_app.dependency_overrides = {}
 
 
 @pytest.fixture(scope="session")
-def engine_chat() -> Generator[Any, None, None]:
-    """
-    Yields a SQLAlchemy engine which is disposed of after the test session
-    """
-    engine = create_engine(DATABASE_URL, echo=True)
-
-    yield engine
-
-    engine.dispose()
-
-
-@pytest.fixture(scope="session")
-def session_chat(engine_chat: Any) -> Generator[Session, None, None]:
+def session_chat(engine: Any) -> Generator[Session, None, None]:
     """
     Yields a SQLAlchemy session within a transaction
     that is rolled back after every session
 
     We need to use the fixture in the session scope because the chat
     endpoint is asynchronous and needs to be open for the entire session
     """
-    connection = engine_chat.connect()
-    # Begin the nested transaction
+    connection = engine.connect()
     transaction = connection.begin()
     # Use connection within the started transaction
     session = Session(bind=connection, query_cls=CustomFilterQuery)
-    # Run Alembic migrations
-    alembic_cfg = Config("src/backend/alembic.ini")
-    upgrade(alembic_cfg, "head")
 
     yield session
 
@@ -123,7 +145,7 @@ def session_chat(engine_chat: Any) -> Generator[Session, None, None]:
 
 
 @pytest.fixture(scope="session")
-def session_client_chat(session_chat: Session) -> Generator[TestClient, None, None]:
+def session_client_chat(session_chat: Session, fastapi_app: FastAPI) -> Generator[TestClient, None, None]:
     """
     Fixture to inject the session into the API client
 
@@ -135,16 +157,14 @@ def session_client_chat(session_chat: Session) -> Generator[TestClient, None, No
     def override_get_session() -> Generator[Session, Any, None]:
         yield session_chat
 
-    app = create_app()
-    app.dependency_overrides[get_session] = override_get_session
+    fastapi_app.dependency_overrides[get_session] = override_get_session
 
     print("Session at fixture " + str(session_chat))
 
-    with TestClient(app) as client:
+    with TestClient(fastapi_app) as client:
         yield client
 
-    app.dependency_overrides = {}
-
+    fastapi_app.dependency_overrides = {}
 
 
 @pytest.fixture(autouse=True)
@@ -158,6 +178,7 @@ def mock_redis_client():
     with patch.object(Redis, 'from_url', return_value=fake_redis):
         yield fake_redis
 
+
 @pytest.fixture
 def user(session: Session) -> User:
     return get_factory("User", session).create(id="1")