Refactor tests, use db migrations, correct session, rollback automatically

backend/pyproject.toml

@@ -20,7 +20,6 @@ pylint = "^2.17.1"
 
 [tool.poetry.group.test.dependencies]
 pytest = "^7.3.1"
-pytest-mock = "^3.10.0"
 requests-mock = "^1.10.0"
 python-multipart = "^0.0.6"
 httpx = "^0.24.0"

backend/src/main.py

@@ -33,7 +33,6 @@
 engine = create_engine(
     "postgresql+pg8000://postgres:password@test-observer-db:5432/postgres", echo=True
 )
-models.Base.metadata.create_all(bind=engine)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 app = FastAPI()
@@ -60,10 +59,8 @@ def root():
 @app.put("/snapmanager")
 def snap_manager(db: Session = Depends(get_db)):
     try:
-        session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
-        with session() as sess:
-            processed_artefacts = snap_manager_controller(sess)
-            logger.info("INFO: Processed artefacts %s", processed_artefacts)
+        processed_artefacts = snap_manager_controller(db)
+        logger.info("INFO: Processed artefacts %s", processed_artefacts)
         if False in processed_artefacts.values():
             return JSONResponse(
                 status_code=500,

backend/src/repository.py

@@ -22,26 +22,6 @@
 from .data_access.models import Family, Stage, Artefact
 
 
-def get_stages_by_family_name(session: Session, family_name: str) -> list | None:
-    """
-    Fetch stages objects related to specific family
-
-    :session: DB session
-    :family_name: name of the family
-    :return: list of stages
-    """
-    family = session.query(Family).filter(Family.name == family_name).first()
-    if family is None:
-        return []
-    stages = (
-        session.query(Stage)
-        .filter(Stage.family_id == family.id)
-        .options(joinedload(Stage.artefacts))
-        .all()
-    )
-    return stages
-
-
 def get_stage_by_name(session: Session, stage_name: str, family: Family) -> Stage:
     """
     Get the stage object by its name
@@ -59,21 +39,9 @@ def get_stage_by_name(session: Session, stage_name: str, family: Family) -> Stag
     return stage
 
 
-def get_family_by_name(session: Session, family_name: str):
-    """
-    Get the family object by its name
-
-    :session: DB session
-    :family_name: Name of the family
-    :return: Family
-    """
-    family = session.query(Family).filter(Family.name == family_name).first()
-    return family
-
-
 def get_artefacts_by_family_name(
     session: Session, family_name: str, is_archived: bool = None
-):
+) -> list[Artefact]:
     """
     Get all the artefacts in a family
 

backend/tests/conftest.py

@@ -16,103 +16,58 @@
 #
 # Written by:
 #        Nadzeya Hutsko <nadzeya.hutsko@canonical.com>
+#        Omar Selo <omar.selo@canonical.com>
 """Fixtures for testing"""
 
 
 import pytest
-from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker, Session
-from sqlalchemy_utils import database_exists, create_database, drop_database
+from alembic import command
+from alembic.config import Config
 from fastapi.testclient import TestClient
-from src.main import app
+from sqlalchemy import Engine, create_engine
+from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy_utils import create_database, database_exists, drop_database
 from src.data_access import Base
-from src.data_access.models import Family, Stage, Artefact
+from src.data_access.models import Artefact, Stage
+from src.main import app, get_db
 
 
-# Setup Test Database
-SQLALCHEMY_DATABASE_URL = (
-    "postgresql+pg8000://postgres:password@test-observer-db:5432/test"
-)
-engine = create_engine(SQLALCHEMY_DATABASE_URL)
-TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+@pytest.fixture(scope="session")
+def db_engine():
+    db_uri = "postgresql+pg8000://postgres:password@test-observer-db:5432/test"
 
+    if not database_exists(db_uri):
+        create_database(db_uri)
 
-@pytest.fixture
-def seed_db(db_session: Session):
-    """Populate database with fake data"""
-    # Snap family
-    family = Family(name="snap")
-    db_session.add(family)
-    # Edge stage
-    stage = Stage(name="edge", family=family, position=10)
-    db_session.add(stage)
-    artefact = Artefact(
-        name="core20", stage=stage, version="1.1.1", source={}, artefact_group=None
-    )
-    db_session.add(artefact)
-    artefact = Artefact(
-        name="docker",
-        stage=stage,
-        version="1.1.1",
-        source={},
-        artefact_group=None,
-        is_archived=True,
-    )
-    db_session.add(artefact)
-    # Beta stage
-    stage = Stage(name="beta", family=family, position=20)
-    db_session.add(stage)
-    artefact = Artefact(
-        name="core22", stage=stage, version="1.1.0", source={}, artefact_group=None
-    )
-    db_session.add(artefact)
+    engine = create_engine(db_uri)
 
-    # Deb family
-    family = Family(name="deb")
-    db_session.add(family)
-    # Proposed stage
-    stage = Stage(name="proposed", family=family, position=10)
-    db_session.add(stage)
-    artefact = Artefact(
-        name="jammy", stage=stage, version="2.1.1", source={}, artefact_group=None
-    )
-    db_session.add(artefact)
-    # Updates stage
-    stage = Stage(name="updates", family=family, position=10)
-    db_session.add(stage)
-    artefact = Artefact(
-        name="raspi", stage=stage, version="2.1.0", source={}, artefact_group=None
-    )
-    db_session.add(artefact)
-    db_session.commit()
+    alembic_config = Config("alembic.ini")
+    alembic_config.set_main_option("sqlalchemy.url", db_uri)
+    command.upgrade(alembic_config, "head")
 
-    yield
+    yield engine
 
-    # Cleanup
-    db_session.query(Artefact).delete()
-    db_session.query(Stage).delete()
-    db_session.query(Family).delete()
-    db_session.commit()
+    Base.metadata.drop_all(engine)
+    engine.dispose()
+    drop_database(db_uri)
 
 
-@pytest.fixture(scope="session")
-def db_session():
-    """Set up and tear down the test database"""
-    if not database_exists(SQLALCHEMY_DATABASE_URL):
-        create_database(SQLALCHEMY_DATABASE_URL)
+@pytest.fixture(scope="function")
+def db_session(db_engine: Engine):
+    connection = db_engine.connect()
+    # Start transaction and not commit it to rollback automatically
+    transaction = connection.begin()
+    session = sessionmaker(autocommit=False, autoflush=False, bind=connection)()
 
-    Base.metadata.create_all(bind=engine)
-    session = TestingSessionLocal()
     yield session
 
-    # Cleanup
     session.close()
-    Base.metadata.drop_all(bind=engine)
-    drop_database(SQLALCHEMY_DATABASE_URL)
+    transaction.close()
+    connection.close()
 
 
-@pytest.fixture(scope="session")
-def test_app():
-    """Create a pytest fixture for the app"""
-    client = TestClient(app)
-    yield client
+@pytest.fixture(scope="function")
+def test_client(db_session: Session) -> TestClient:
+    """Create a test http client"""
+    app.dependency_overrides[get_db] = lambda: db_session
+    return TestClient(app)

backend/tests/helpers.py

@@ -0,0 +1,18 @@
+from sqlalchemy.orm import Session
+from src.data_access.models import Artefact, Stage
+
+
+def create_artefact(db_session: Session, stage_name: str, **kwargs):
+    """Create a dummy artefact"""
+    stage = db_session.query(Stage).filter(Stage.name == stage_name).first()
+    artefact = Artefact(
+        name=kwargs.get("name", ""),
+        stage=stage,
+        version=kwargs.get("version", "1.1.1"),
+        source=kwargs.get("source", {}),
+        artefact_group=None,
+        is_archived=kwargs.get("is_archived", False),
+    )
+    db_session.add(artefact)
+    db_session.commit()
+    return artefact