Тестирование SQLAlchemy моделей, миграций, транзакций
«Нахождение багов в работе с данными до попадания в production.»
from sqlalchemy import Column, Integer, String, Decimal, ForeignKey, create_engine
from sqlalchemy.orm import sessionmaker, relationship, declarative_base
from hypothesis import given, strategies as st
import pytest
Base = declarative_base()
class User(Base):
__tablename__ = 'users'
id = Column(Integer, primary_key=True)
email = Column(String, unique=True, nullable=False)
balance = Column(Decimal, default=0)
orders = relationship('Order', back_populates='user')
class Order(Base):
__tablename__ = 'orders'
id = Column(Integer, primary_key=True)
user_id = Column(Integer, ForeignKey('users.id'))
total = Column(Decimal)
user = relationship('User', back_populates='orders')
# Фикстура для тестовой БД
@pytest.fixture
def session():
engine = create_engine('sqlite:///:memory:')
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)
session = Session()
yield session
session.rollback()
session.close()@st.composite
def user_strategy(draw):
return {
'email': draw(st.emails()),
'balance': draw(st.decimals(min_value=0, max_value=1000000, places=2))
}
@given(user_strategy())
def test_create_user(session, user_data):
user = User(**user_data)
session.add(user)
session.commit()
assert user.id is not None
assert user.email == user_data['email']
assert user.balance == user_data['balance']@st.composite
def order_with_user(draw):
user_data = draw(user_strategy())
order_total = draw(st.decimals(min_value=0.01, max_value=100000, places=2))
return {
'user_data': user_data,
'order_total': order_total
}
@given(order_with_user())
def test_user_order_relationship(session, data):
user = User(**data['user_data'])
order = Order(total=data['order_total'])
user.orders.append(order)
session.add(user)
session.commit()
assert len(user.orders) == 1
assert user.orders[0].total == data['order_total']
assert order.user.id == user.idfrom contextlib import contextmanager
@contextmanager
def transaction(session):
try:
yield session
session.commit()
except:
session.rollback()
raise
@given(st.decimals(min_value=0.01, max_value=10000, places=2))
def test_transfer_atomicity(session, amount):
"""Транзакция перевода должна быть атомарной"""
user1 = User(email='user1@test.com', balance=10000)
user2 = User(email='user2@test.com', balance=0)
session.add_all([user1, user2])
session.commit()
try:
with transaction(session):
user1.balance -= amount
user2.balance += amount
except:
pass # Ожидаем возможный rollback
# Инвариант: общая сумма не изменилась
session.refresh(user1)
session.refresh(user2)
assert user1.balance + user2.balance == 10000from sqlalchemy.exc import IntegrityError
@given(st.emails())
def test_unique_email_constraint(session, email):
"""Проверка unique constraint на email"""
user1 = User(email=email)
user2 = User(email=email) # Тот же email
session.add(user1)
session.commit()
session.add(user2)
with pytest.raises(IntegrityError):
session.commit()from alembic.command import upgrade, downgrade
from alembic.config import Config
def test_migration_upgrade_downgrade(session):
"""Проверка что миграции работают в обе стороны"""
alembic_cfg = Config('alembic.ini')
# Upgrade до последней
upgrade(alembic_cfg, 'head')
# Проверка что данные сохранились
users = session.query(User).all()
original_count = len(users)
# Downgrade на одну миграцию назад
downgrade(alembic_cfg, '-1')
# Upgrade снова
upgrade(alembic_cfg, 'head')
# Данные должны сохраниться
users_after = session.query(User).all()
assert len(users_after) == original_countHypothesis автоматически находит баги в ORM маппинге, транзакциях, constraints и миграциях.
Следующая тема: Валидация данных — Pydantic, JSON-схемы, сериализация.
Далее: Валидация данных и схем