88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
"""Real PostgreSQL tests, isolated in a fresh schema for every test."""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from fastapi.testclient import TestClient
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.engine import make_url
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from app.db.session import get_db
|
|
from app.main import create_app
|
|
|
|
BACKEND = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def migrate(engine, revision="head", *, downgrade=False):
|
|
config = Config(str(BACKEND / "alembic.ini"))
|
|
with engine.begin() as connection:
|
|
config.attributes["connection"] = connection
|
|
(command.downgrade if downgrade else command.upgrade)(config, revision)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def postgres_url(tmp_path_factory):
|
|
url = os.environ.get("TEST_DATABASE_URL")
|
|
server = None
|
|
if not url:
|
|
try:
|
|
import pgserver
|
|
except ImportError:
|
|
pytest.fail("请设置 TEST_DATABASE_URL，或安装可选的本机测试依赖 pgserver")
|
|
server = pgserver.get_server(tmp_path_factory.mktemp("postgres"), cleanup_mode="stop")
|
|
url = server.get_uri().replace("postgresql://", "postgresql+psycopg://", 1)
|
|
if not url.startswith("postgresql+psycopg://"):
|
|
pytest.fail("TEST_DATABASE_URL 必须使用 postgresql+psycopg://")
|
|
bootstrap = create_engine(url)
|
|
with bootstrap.begin() as connection:
|
|
connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector WITH SCHEMA public"))
|
|
bootstrap.dispose()
|
|
yield url
|
|
if server:
|
|
server.cleanup()
|
|
|
|
|
|
@pytest.fixture
|
|
def database(postgres_url):
|
|
schema = "test_" + uuid4().hex
|
|
admin = create_engine(postgres_url)
|
|
with admin.begin() as connection:
|
|
connection.execute(text(f'CREATE SCHEMA "{schema}"'))
|
|
# 不包含 public，防止误操作已有业务表。
|
|
url = make_url(postgres_url).update_query_dict({"options": f"-csearch_path={schema}"})
|
|
engine = create_engine(url)
|
|
yield engine
|
|
engine.dispose()
|
|
with admin.begin() as connection:
|
|
connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))
|
|
admin.dispose()
|
|
|
|
|
|
@pytest.fixture
|
|
def migrated_database(database):
|
|
migrate(database)
|
|
return database
|
|
|
|
|
|
@pytest.fixture
|
|
def client(migrated_database):
|
|
app = create_app(migrated_database)
|
|
factory = sessionmaker(bind=migrated_database)
|
|
|
|
def test_session():
|
|
with factory() as session:
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
|
|
app.dependency_overrides[get_db] = test_session
|
|
with TestClient(app) as client:
|
|
yield client
|