100 lines
2.6 KiB
Python
100 lines
2.6 KiB
Python
"""
|
|
The functions in this module can be used for testing that the constraints of
|
|
your models. Each assert function runs SQL UPDATEs that check for the existence
|
|
of given constraint. Consider the following model::
|
|
|
|
|
|
class User(Base):
|
|
__tablename__ = 'user'
|
|
id = sa.Column(sa.Integer, primary_key=True)
|
|
name = sa.Column(sa.String(200), nullable=True)
|
|
email = sa.Column(sa.String(255), nullable=False)
|
|
|
|
|
|
user = User(name='John Doe', email='john@example.com')
|
|
session.add(user)
|
|
session.commit()
|
|
|
|
|
|
We can easily test the constraints by assert_* functions::
|
|
|
|
|
|
from sqlalchemy_utils import (
|
|
assert_nullable,
|
|
assert_non_nullable,
|
|
assert_max_length
|
|
)
|
|
|
|
assert_nullable(user, 'name')
|
|
assert_non_nullable(user, 'email')
|
|
assert_max_length(user, 'name', 200)
|
|
|
|
# raises AssertionError because the max length of email is 255
|
|
assert_max_length(user, 'email', 300)
|
|
"""
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.exc import DataError, IntegrityError
|
|
|
|
|
|
class raises(object):
|
|
def __init__(self, expected_exc):
|
|
self.expected_exc = expected_exc
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if exc_type != self.expected_exc:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _update_field(obj, field, value):
|
|
session = sa.orm.object_session(obj)
|
|
table = sa.inspect(obj.__class__).columns[field].table
|
|
query = table.update().values(**{field: value})
|
|
session.execute(query)
|
|
session.flush()
|
|
|
|
|
|
def assert_nullable(obj, column):
|
|
"""
|
|
Assert that given column is nullable. This is checked by running an SQL
|
|
update that assigns given column as None.
|
|
|
|
:param obj: SQLAlchemy declarative model object
|
|
:param column: Name of the column
|
|
"""
|
|
try:
|
|
_update_field(obj, column, None)
|
|
except (IntegrityError) as e:
|
|
assert False, str(e)
|
|
|
|
|
|
def assert_non_nullable(obj, column):
|
|
"""
|
|
Assert that given column is not nullable. This is checked by running an SQL
|
|
update that assigns given column as None.
|
|
|
|
:param obj: SQLAlchemy declarative model object
|
|
:param column: Name of the column
|
|
"""
|
|
with raises(IntegrityError):
|
|
_update_field(obj, column, None)
|
|
|
|
|
|
def assert_max_length(obj, column, max_length):
|
|
"""
|
|
Assert that the given column is of given max length.
|
|
|
|
:param obj: SQLAlchemy declarative model object
|
|
:param column: Name of the column
|
|
"""
|
|
try:
|
|
_update_field(obj, column, u'a' * max_length)
|
|
except (DataError) as e:
|
|
assert False, str(e)
|
|
with raises(DataError):
|
|
_update_field(obj, column, u'a' * (max_length + 1))
|
|
|