from alembic.config import Config
from alembic import command
import pytest
from sqlalchemy import create_engine, text
from sqlalchemy.pool import NullPool
from sqlalchemy.orm import Session
from sqlalchemy.engine import Engine
import os
import random
from typing import Generator
from pgtemp_example_py import insert_task, list_tasks, complete_task
def get_db_and_run_migrations() -> tuple[Engine, str]:
db_url = os.environ["PGTEMP_SINGLE"]
engine = create_engine(db_url)
db_id = random.randint(0, 1000000)
db_name = f"test{db_id}"
with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
conn.execute(text(f"CREATE DATABASE {db_name};"))
print(f"database name: `{db_name}`")
db_url += "/" + db_name
cfg = Config("alembic.ini")
cfg.set_main_option("sqlalchemy.url", db_url)
command.upgrade(cfg, "head")
engine = create_engine(db_url, poolclass=NullPool)
return engine, db_name
def drop_db(db_name: str):
db_url = os.environ["PGTEMP_SINGLE"]
engine = create_engine(db_url)
with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
conn.execute(text(f"DROP DATABASE {db_name}"))
@pytest.fixture
def test_db() -> Generator[Engine, None, None]:
db, db_name = get_db_and_run_migrations()
yield db
db.dispose()
drop_db(db_name)
def test_basic_ops(test_db):
with Session(test_db) as session:
insert_task(session, "hello")
session.commit()
with Session(test_db) as session:
tasks = list_tasks(session)
assert len(tasks) == 1
t = tasks[0]
tid = tasks[0].id
assert t.task == "hello"
assert t.completed is False
session.commit()
with Session(test_db) as session:
complete_task(session, tid)
session.commit()
with Session(test_db) as session:
tasks = list_tasks(session)
assert len(tasks) == 1
t = tasks[0]
assert t.task == "hello"
assert t.completed is True
session.commit()