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/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

backend/tests/test_repository.py

@@ -16,24 +16,26 @@
 #
 # Written by:
 #        Nadzeya Hutsko <nadzeya.hutsko@canonical.com>
+#        Omar Selo <omar.selo@canonical.com>
 """Test services functions"""
 
 
-import pytest
-
+from sqlalchemy.orm import Session
 from src.repository import (
-    get_stages_by_family_name,
-    get_stage_by_name,
-    get_family_by_name,
     get_artefacts_by_family_name,
+    get_family_by_name,
+    get_stage_by_name,
+    get_stages_by_family_name,
 )
 
+from .helpers import create_artefact
 
-def test_get_stages_by_family_name(db_session, seed_db):
+
+def test_get_stages_by_family_name(db_session: Session):
     """The function should select correct stages for the specified family name"""
     # Arrange
     family_name = "snap"
-    expected_stage_names = ["edge", "beta"]
+    expected_stage_names = ["edge", "beta", "candidate", "stable"]
 
     # Act
     stages = get_stages_by_family_name(db_session, family_name)
@@ -43,7 +45,7 @@ def test_get_stages_by_family_name(db_session, seed_db):
     assert all(stage.name in expected_stage_names for stage in stages)
 
 
-def test_get_stages_by_family_name_no_such_family(seed_db, db_session):
+def test_get_stages_by_family_name_no_such_family(db_session: Session):
     """The function should return empty list"""
     # Arrange
     family_name = "fake"
@@ -55,7 +57,7 @@ def test_get_stages_by_family_name_no_such_family(seed_db, db_session):
     assert stages == []
 
 
-def test_get_stage_by_name(seed_db, db_session):
+def test_get_stage_by_name(db_session: Session):
     """The function should select the correct stage by its name"""
     # Arrange
     family = get_family_by_name(db_session, "deb")
@@ -68,7 +70,7 @@ def test_get_stage_by_name(seed_db, db_session):
     assert stage.name == stage_name
 
 
-def test_get_stage_by_name_no_such_stage(seed_db, db_session):
+def test_get_stage_by_name_no_such_stage(db_session: Session):
     """The function should return None"""
     # Arrange
     family = get_family_by_name(db_session, "deb")
@@ -81,49 +83,36 @@ def test_get_stage_by_name_no_such_stage(seed_db, db_session):
     assert stage is None
 
 
-def test_get_stage_by_name_no_such_family(seed_db, db_session):
-    """The function should return None"""
-    # Arrange
-    family = get_family_by_name(db_session, "deb")
-    stage_name = "fakestage"
-
-    # Act
-    stage = get_stage_by_name(db_session, stage_name, family)
-
-    # Assert
-    assert stage is None
-
-
-def test_get_artefacts_by_family_name(seed_db, db_session):
+def test_get_artefacts_by_family_name(db_session: Session):
     """We should get a valid list of artefacts"""
     # Arrange
-    family_name = "snap"
     expected_artefact_names = ["core20", "core22", "docker"]
+    for name in expected_artefact_names:
+        create_artefact(db_session, "beta", name=name)
 
     # Act
-    artefacts = get_artefacts_by_family_name(db_session, family_name)
+    artefacts = get_artefacts_by_family_name(db_session, "snap")
 
     # Assert
     assert len(artefacts) == len(expected_artefact_names)
     assert all(artefact.name in expected_artefact_names for artefact in artefacts)
 
 
-def test_get_artefacts_by_family_name_filter_archived(seed_db, db_session):
+def test_get_artefacts_by_family_name_filter_archived(db_session: Session):
     """We should get a list of archived artefacts"""
     # Arrange
-    family_name = "snap"
-    expected_artefact_names = ["docker"]
+    create_artefact(db_session, "beta", name="docker", is_archived=True)
 
     # Act
-    artefacts = get_artefacts_by_family_name(db_session, family_name, is_archived=True)
+    artefacts = get_artefacts_by_family_name(db_session, "snap", is_archived=True)
 
     # Assert
-    assert len(artefacts) == len(expected_artefact_names)
-    assert all(artefact.name in expected_artefact_names for artefact in artefacts)
-    assert all(artefact.is_archived for artefact in artefacts)
+    assert len(artefacts) == 1
+    assert artefacts[0].name == "docker"
+    assert artefacts[0].is_archived
 
 
-def test_get_artefacts_by_family_name_no_such_family(seed_db, db_session):
+def test_get_artefacts_by_family_name_no_such_family(db_session: Session):
     """We should get an empty list when there's no such family"""
     # Arrange
     family_name = "fakename"