|
16 | 16 | #
|
17 | 17 | # Written by:
|
18 | 18 | # Nadzeya Hutsko <nadzeya.hutsko@canonical.com>
|
| 19 | +# Omar Selo <omar.selo@canonical.com> |
19 | 20 | """Fixtures for testing"""
|
20 | 21 |
|
21 | 22 |
|
22 | 23 | import pytest
|
23 |
| -from sqlalchemy import create_engine |
24 |
| -from sqlalchemy.orm import sessionmaker, Session |
25 |
| -from sqlalchemy_utils import database_exists, create_database, drop_database |
| 24 | +from alembic import command |
| 25 | +from alembic.config import Config |
26 | 26 | from fastapi.testclient import TestClient
|
27 |
| -from src.main import app |
| 27 | +from sqlalchemy import Engine, create_engine |
| 28 | +from sqlalchemy.orm import Session, sessionmaker |
| 29 | +from sqlalchemy_utils import create_database, database_exists, drop_database |
28 | 30 | from src.data_access import Base
|
29 |
| -from src.data_access.models import Family, Stage, Artefact |
| 31 | +from src.data_access.models import Artefact, Stage |
| 32 | +from src.main import app, get_db |
30 | 33 |
|
31 | 34 |
|
32 |
| -# Setup Test Database |
33 |
| -SQLALCHEMY_DATABASE_URL = ( |
34 |
| - "postgresql+pg8000://postgres:password@test-observer-db:5432/test" |
35 |
| -) |
36 |
| -engine = create_engine(SQLALCHEMY_DATABASE_URL) |
37 |
| -TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) |
| 35 | +@pytest.fixture(scope="session") |
| 36 | +def db_engine(): |
| 37 | + db_uri = "postgresql+pg8000://postgres:password@test-observer-db:5432/test" |
38 | 38 |
|
| 39 | + if not database_exists(db_uri): |
| 40 | + create_database(db_uri) |
39 | 41 |
|
40 |
| -@pytest.fixture |
41 |
| -def seed_db(db_session: Session): |
42 |
| - """Populate database with fake data""" |
43 |
| - # Snap family |
44 |
| - family = Family(name="snap") |
45 |
| - db_session.add(family) |
46 |
| - # Edge stage |
47 |
| - stage = Stage(name="edge", family=family, position=10) |
48 |
| - db_session.add(stage) |
49 |
| - artefact = Artefact( |
50 |
| - name="core20", stage=stage, version="1.1.1", source={}, artefact_group=None |
51 |
| - ) |
52 |
| - db_session.add(artefact) |
53 |
| - artefact = Artefact( |
54 |
| - name="docker", |
55 |
| - stage=stage, |
56 |
| - version="1.1.1", |
57 |
| - source={}, |
58 |
| - artefact_group=None, |
59 |
| - is_archived=True, |
60 |
| - ) |
61 |
| - db_session.add(artefact) |
62 |
| - # Beta stage |
63 |
| - stage = Stage(name="beta", family=family, position=20) |
64 |
| - db_session.add(stage) |
65 |
| - artefact = Artefact( |
66 |
| - name="core22", stage=stage, version="1.1.0", source={}, artefact_group=None |
67 |
| - ) |
68 |
| - db_session.add(artefact) |
| 42 | + engine = create_engine(db_uri) |
69 | 43 |
|
70 |
| - # Deb family |
71 |
| - family = Family(name="deb") |
72 |
| - db_session.add(family) |
73 |
| - # Proposed stage |
74 |
| - stage = Stage(name="proposed", family=family, position=10) |
75 |
| - db_session.add(stage) |
76 |
| - artefact = Artefact( |
77 |
| - name="jammy", stage=stage, version="2.1.1", source={}, artefact_group=None |
78 |
| - ) |
79 |
| - db_session.add(artefact) |
80 |
| - # Updates stage |
81 |
| - stage = Stage(name="updates", family=family, position=10) |
82 |
| - db_session.add(stage) |
83 |
| - artefact = Artefact( |
84 |
| - name="raspi", stage=stage, version="2.1.0", source={}, artefact_group=None |
85 |
| - ) |
86 |
| - db_session.add(artefact) |
87 |
| - db_session.commit() |
| 44 | + alembic_config = Config("alembic.ini") |
| 45 | + alembic_config.set_main_option("sqlalchemy.url", db_uri) |
| 46 | + command.upgrade(alembic_config, "head") |
88 | 47 |
|
89 |
| - yield |
| 48 | + yield engine |
90 | 49 |
|
91 |
| - # Cleanup |
92 |
| - db_session.query(Artefact).delete() |
93 |
| - db_session.query(Stage).delete() |
94 |
| - db_session.query(Family).delete() |
95 |
| - db_session.commit() |
| 50 | + Base.metadata.drop_all(engine) |
| 51 | + engine.dispose() |
| 52 | + drop_database(db_uri) |
96 | 53 |
|
97 | 54 |
|
98 |
| -@pytest.fixture(scope="session") |
99 |
| -def db_session(): |
100 |
| - """Set up and tear down the test database""" |
101 |
| - if not database_exists(SQLALCHEMY_DATABASE_URL): |
102 |
| - create_database(SQLALCHEMY_DATABASE_URL) |
| 55 | +@pytest.fixture(scope="function") |
| 56 | +def db_session(db_engine: Engine): |
| 57 | + connection = db_engine.connect() |
| 58 | + # Start transaction and not commit it to rollback automatically |
| 59 | + transaction = connection.begin() |
| 60 | + session = sessionmaker(autocommit=False, autoflush=False, bind=connection)() |
103 | 61 |
|
104 |
| - Base.metadata.create_all(bind=engine) |
105 |
| - session = TestingSessionLocal() |
106 | 62 | yield session
|
107 | 63 |
|
108 |
| - # Cleanup |
109 | 64 | session.close()
|
110 |
| - Base.metadata.drop_all(bind=engine) |
111 |
| - drop_database(SQLALCHEMY_DATABASE_URL) |
| 65 | + transaction.close() |
| 66 | + connection.close() |
112 | 67 |
|
113 | 68 |
|
114 |
| -@pytest.fixture(scope="session") |
115 |
| -def test_app(): |
116 |
| - """Create a pytest fixture for the app""" |
117 |
| - client = TestClient(app) |
118 |
| - yield client |
| 69 | +@pytest.fixture(scope="function") |
| 70 | +def test_client(db_session: Session) -> TestClient: |
| 71 | + """Create a test http client""" |
| 72 | + app.dependency_overrides[get_db] = lambda: db_session |
| 73 | + return TestClient(app) |
0 commit comments