#!/bin/python

import os
from collections.abc import Generator

os.environ["DATABASE_URL"] = "sqlite://"

import pytest
from alembic.config import Config
from sqlalchemy import StaticPool, create_engine
from sqlalchemy.orm import Session, sessionmaker

from alembic import command
from src.db import set_engine


@pytest.fixture()
def db_session() -> Generator[Session]:
    engine = create_engine(
        "sqlite://",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )

    set_engine(engine)

    alembic_cfg = Config("alembic.ini")
    command.upgrade(alembic_cfg, "head")

    SessionLocal = sessionmaker(bind=engine)
    session = SessionLocal()
    try:
        yield session
    finally:
        session.close()
        engine.dispose()