src/tests/conftest.py - master - urusai/code_sample

src/tests/conftest.py
68 строк2 KB
nick
linters
30 сен 2024, 16:13
30 сен 2024, 16:1309a0f4a
import asyncio

import pytest
import fastapi
import pytest_asyncio
from models import base
from core.settings import settings
from sqlalchemy.pool import NullPool
from sqlalchemy_utils import drop_database, create_database, database_exists
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from main import app as default_app

TEST_SYNC_DATABASE_URL = settings.TEST_SYNC_DATABASE_URL
TEST_ASYNC_DATABASE_URL = settings.TEST_ASYNC_DATABASE_URL
async_engine = create_async_engine(TEST_ASYNC_DATABASE_URL, poolclass=NullPool)
async_session = async_sessionmaker(
    bind=async_engine, expire_on_commit=False, autoflush=False
)


class FastAPI(fastapi.FastAPI):
    # чтобы убрать предупреждение об отсутствии атрибута (это не правда)
    dependency_overrides: dict


def create_db():
    if not database_exists(TEST_SYNC_DATABASE_URL):
        create_database(TEST_SYNC_DATABASE_URL)


async def create_tables():
    async with async_engine.begin() as conn:
        await conn.run_sync(base.BaseModel.metadata.create_all)


def drop_db():
    if database_exists(TEST_SYNC_DATABASE_URL):
        drop_database(TEST_SYNC_DATABASE_URL)


@pytest.fixture(scope="session", autouse=True)
def init_database():
    create_db()
    asyncio.run(create_tables())
    yield
    drop_db()


@pytest_asyncio.fixture(scope="function", autouse=True)
async def db_connection(init_database):
    connection = await async_engine.connect()
    transaction = await connection.begin()
    yield connection
    await transaction.rollback()
    await connection.close()


@pytest_asyncio.fixture(scope="function")
async def session(db_connection):
    session = AsyncSession(bind=db_connection)
    return session


@pytest.fixture(scope="function")
def app() -> FastAPI:
    yield default_app
    default_app.dependency_overrides = {}