Conftest to create Postgres Database and auto delete data after test end

conftest.py

import os

import pytest
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy_utils.functions import create_database, database_exists, drop_database
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker

from main import app as main_app
from main.database import db as _db

@pytest.fixture(scope="session")
def database(request):
    main_app.config.from_object("main.config.Test")
    sqlalchemy_database_uri = main_app.config.get("SQLALCHEMY_DATABASE_URI")

    # It only happens on the first run. Needs further investigation.
    if database_exists(sqlalchemy_database_uri):
        drop_database(sqlalchemy_database_uri)

    create_database(sqlalchemy_database_uri)

    @request.addfinalizer
    def drop_databases():
        drop_database(sqlalchemy_database_uri)


@pytest.fixture(scope="session")
def app(database):
    main_app.config.from_object("main.config.Test")
    yield main_app


@pytest.fixture(scope="session")
def client(app):
    client = app.test_client()
    yield client


@pytest.fixture(scope="session")
def db_engine(app):
    """yields a SQLAlchemy engine which is suppressed after the test session"""
    _db.app = app
    _db.create_all()

    yield _db


@pytest.fixture(scope="function", autouse=True)
def db_session(db_engine):
    """yields a SQLAlchemy connection which is rollbacked after the test"""
    connection = db_engine.engine.connect()
    transaction = connection.begin()

    options = dict(bind=connection, binds={})
    session_ = db_engine.create_scoped_session(options=options)

    db_engine.session = session_

    yield session_

    transaction.rollback()
    connection.close()
    session_.remove()