120 lines
3.8 KiB
Python
120 lines
3.8 KiB
Python
import os
|
|
from collections.abc import Callable, Generator
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
os.environ["DATABASE_URL"] = "sqlite://"
|
|
|
|
import pytest
|
|
from alembic.config import Config
|
|
from sqlalchemy import Engine, StaticPool, create_engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
from alembic import command
|
|
from src.db import set_engine
|
|
from src.ingest.b4w import upsert_b4w
|
|
from tests.factories import B4wExportFactory
|
|
|
|
DEFAULT_B4W_ROWS = 10
|
|
# ``*.db`` and ``.local`` are gitignored, so artifacts are never committed.
|
|
TEST_DB_DIR = Path(os.environ.get("TEST_DB_DIR", ".local/test-dbs"))
|
|
|
|
|
|
@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()
|
|
|
|
|
|
@pytest.fixture()
|
|
def db_session_factory() -> Generator[Callable[[str], Session]]:
|
|
"""Create named, file-backed SQLite databases that survive the test run.
|
|
|
|
Each call recreates ``<TEST_DB_DIR>/<name>.db`` from scratch (overwriting
|
|
the artifact of the previous run), runs the migrations and returns a
|
|
session. The files are intentionally *not* removed on teardown, so the last
|
|
created instance can be inspected afterwards, e.g.::
|
|
|
|
sqlite3 .local/test-dbs/b4w_basic.db
|
|
"""
|
|
engines: list[Engine] = []
|
|
|
|
def make(name: str) -> Session:
|
|
TEST_DB_DIR.mkdir(parents=True, exist_ok=True)
|
|
db_path = (TEST_DB_DIR / f"{name}.db").resolve()
|
|
# Start from a clean database: overwrite whatever the last run left.
|
|
for suffix in ("", "-wal", "-shm"):
|
|
Path(f"{db_path}{suffix}").unlink(missing_ok=True)
|
|
|
|
engine = create_engine(f"sqlite:///{db_path}")
|
|
engines.append(engine)
|
|
set_engine(engine) # alembic/env.py picks this up via get_set_engine()
|
|
|
|
alembic_cfg = Config("alembic.ini")
|
|
command.upgrade(alembic_cfg, "head")
|
|
|
|
return sessionmaker(bind=engine)()
|
|
|
|
try:
|
|
yield make
|
|
finally:
|
|
for engine in engines:
|
|
engine.dispose()
|
|
|
|
|
|
@pytest.fixture()
|
|
def populated_db_b4w(
|
|
db_session_factory: Callable[[str], Session],
|
|
) -> Callable[..., Session]:
|
|
"""Populate ``import_b4w`` (and ``category``) in a named, persisted DB.
|
|
|
|
Only the first-layer ingestion table is filled; no further processing.
|
|
``**kwargs`` are forwarded to ``B4wExportFactory.build``; callable values
|
|
are resolved per row (e.g. ``ValDt=lambda: faker.date_between(...)``).
|
|
An optional ``setup`` callback runs after database creation and before rows
|
|
are inserted.
|
|
``Id1`` is reserved: it is set per row (``1..row_count``) and cannot be
|
|
overridden via ``**kwargs``.
|
|
The returned session is *not* closed by the fixture; close it in the test.
|
|
"""
|
|
|
|
def populate_db_b4w(
|
|
name: str,
|
|
row_count: int = DEFAULT_B4W_ROWS,
|
|
*,
|
|
setup: Callable[[Session], None] | None = None,
|
|
**kwargs: Any,
|
|
) -> Session:
|
|
if "Id1" in kwargs:
|
|
raise ValueError(
|
|
"Id1 is set per row (1..row_count) by populated_db_b4w and "
|
|
"cannot be overridden via **kwargs"
|
|
)
|
|
session = db_session_factory(name)
|
|
if setup is not None:
|
|
setup(session)
|
|
for i in range(row_count):
|
|
row_kwargs = {k: (v() if callable(v) else v) for k, v in kwargs.items()}
|
|
row = B4wExportFactory.build(Id1=i + 1, **cast(dict[str, Any], row_kwargs))
|
|
upsert_b4w(session, row, do_commit=False)
|
|
session.commit()
|
|
return session
|
|
|
|
return populate_db_b4w
|