110 lines
3.4 KiB
Python
110 lines
3.4 KiB
Python
"""PostgreSQL 后端集成测试（需 CHONGMING_TEST_PG_URL 环境变量，否则跳过）。"""
|
|
|
|
import os
|
|
import uuid
|
|
|
|
import asyncpg
|
|
import pytest
|
|
from sqlmodel import Field, SQLModel, select
|
|
|
|
from chongming_database import Database
|
|
|
|
PG_URL = os.environ.get("CHONGMING_TEST_PG_URL")
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not PG_URL, reason="CHONGMING_TEST_PG_URL 未设置，跳过 PostgreSQL 集成测试"
|
|
)
|
|
|
|
|
|
def _make_item_table(tablename: str):
|
|
"""动态生成随机表名的 SQLModel 表，避免污染共享 PG 实例。"""
|
|
|
|
class Item(SQLModel, table=True):
|
|
__tablename__ = tablename
|
|
|
|
id: int | None = Field(default=None, primary_key=True)
|
|
name: str = Field(max_length=50)
|
|
|
|
return Item
|
|
|
|
|
|
async def test_postgres_crud():
|
|
Item = _make_item_table(f"pg_test_item_{uuid.uuid4().hex[:8]}")
|
|
db = Database.from_url(PG_URL)
|
|
await db.initialize()
|
|
try:
|
|
async with db.get_session() as session:
|
|
session.add(Item(name="pg"))
|
|
await session.commit()
|
|
async with db.get_session() as session:
|
|
rows = (await session.execute(select(Item))).scalars().all()
|
|
assert [r.name for r in rows] == ["pg"]
|
|
finally:
|
|
await db.dispose()
|
|
|
|
|
|
async def test_postgres_bare_url_scheme():
|
|
"""裸 postgresql:// 应被归一化为 postgresql+asyncpg:// 后正常连接。"""
|
|
bare = PG_URL.replace("postgresql+asyncpg://", "postgresql://", 1)
|
|
db = Database.from_url(bare)
|
|
await db.initialize()
|
|
try:
|
|
async with db.get_session() as session:
|
|
await session.execute(select(1))
|
|
finally:
|
|
await db.dispose()
|
|
|
|
|
|
async def test_from_config_postgres_end_to_end():
|
|
db = Database.from_config({"type": "postgresql", "url": PG_URL})
|
|
await db.initialize()
|
|
try:
|
|
async with db.get_session() as session:
|
|
await session.execute(select(1))
|
|
finally:
|
|
await db.dispose()
|
|
|
|
|
|
async def test_postgres_generate_schemas_false():
|
|
db = Database.from_url(PG_URL)
|
|
await db.initialize(generate_schemas=False)
|
|
try:
|
|
async with db.get_session() as session:
|
|
await session.execute(select(1))
|
|
finally:
|
|
await db.dispose()
|
|
|
|
|
|
async def test_postgres_auto_create_missing_database():
|
|
"""目标库不存在时 initialize() 应自动建库（回归：database "x" does not exist）。"""
|
|
from sqlalchemy.engine import make_url
|
|
|
|
base = make_url(PG_URL)
|
|
db_name = f"chongming_test_autocreate_{uuid.uuid4().hex[:8]}"
|
|
url = base.set(database=db_name).render_as_string(hide_password=False)
|
|
|
|
db = Database.from_url(url)
|
|
try:
|
|
await db.initialize() # 此前这里会因库不存在而抛 OperationalError
|
|
async with db.get_session() as session:
|
|
cur = await session.execute(select(1))
|
|
assert cur.scalar() == 1
|
|
# 幂等：库已存在时重复 initialize 也应成功（先释放旧引擎连接，便于清理）
|
|
await db.dispose()
|
|
await db.initialize()
|
|
finally:
|
|
await db.dispose()
|
|
# 清理：连 postgres 维护库删除测试库
|
|
maint = base.set(database="postgres")
|
|
conn = await asyncpg.connect(
|
|
host=maint.host,
|
|
port=maint.port,
|
|
user=maint.username,
|
|
password=maint.password,
|
|
database="postgres",
|
|
)
|
|
try:
|
|
await conn.execute(f'DROP DATABASE IF EXISTS "{db_name}"')
|
|
finally:
|
|
await conn.close()
|