Test SQLAlchemy use of PostgreSQL ON DELETE CASCADE

test-psql-ondelete-cascade.py

import unittest

from sqlalchemy import create_engine
from sqlalchemy import Column, Integer
from sqlalchemy import String
from sqlalchemy import Text
from sqlalchemy import ForeignKey
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import scoped_session
from sqlalchemy.orm import relationship
from sqlalchemy.orm import backref
from sqlalchemy.ext.declarative import declarative_base


class PostgresOnDeleteCascadeTest(unittest.TestCase):

    def setUp(self):
        Session = scoped_session(sessionmaker(expire_on_commit=False))
        class Model(object):
            query = Session.query_property()

            def save(self, session=Session):
                session.add(self)
                session.commit()
                return self

            def delete(self, session=Session):
                session.delete(self)
                session.commit()

        Base = declarative_base(cls=Model)

        class User(Base):
            __tablename__ = 'users'

            id = Column(Integer, primary_key=True)
            username = Column(String(255), unique=True, nullable=False)

        engine = create_engine('postgresql://test@localhost/test', echo=True)
        Session.configure(bind=engine)
        Base.metadata.bind = engine
        User.__table__.create()
        self.engine = engine
        self.Base = Base
        self.Session = Session
        self.User = User

    def tearDown(self):
        Session = self.Session
        Base = self.Base
        Session.expunge_all()
        Session.rollback()
        Session.remove()
        Base.metadata.drop_all()

    # Passes
    def test_cascade_delete_using_sql(self):
        Base = self.Base
        User = self.User
        engine = self.engine
        class Post(Base):
            __tablename__ = 'posts'

            id = Column(Integer, primary_key=True)
            title = Column(String(255), nullable=False)
            body = Column(Text, nullable=False)
            user_id = Column(Integer, ForeignKey(User.id, ondelete='CASCADE'),
                             nullable=False)
            user = relationship(User, backref='posts')

        Post.__table__.create()
        user = User(username='dude').save()
        post = Post(user=user, title='foo', body='bar').save()
        engine.execute('DELETE FROM users WHERE id = %d' % user.id)
        self.assertEqual(User.query.count(), 0)
        self.assertEqual(Post.query.count(), 0)

    # Error because it tries to set post.user_id = None
    def test_cascade_delete_posts_null_constraint(self):
        Base = self.Base
        User = self.User
        class Post(Base):
            __tablename__ = 'posts'

            id = Column(Integer, primary_key=True)
            title = Column(String(255), nullable=False)
            body = Column(Text, nullable=False)
            user_id = Column(Integer, ForeignKey(User.id, ondelete='CASCADE'),
                             nullable=False)
            user = relationship(User, backref='posts')

        Post.__table__.create()
        user = User(username='dude').save()
        post = Post(user=user, title='foo', body='bar').save()
        user.delete()
        self.assertEqual(User.query.count(), 0)
        self.assertEqual(Post.query.count(), 0)

    # Fails to delete post
    def test_cascade_delete_posts(self):
        Base = self.Base
        User = self.User
        class Post(Base):
            __tablename__ = 'posts'

            id = Column(Integer, primary_key=True)
            title = Column(String(255), nullable=False)
            body = Column(Text, nullable=False)
            user_id = Column(Integer, ForeignKey(User.id, ondelete='CASCADE'))
            user = relationship(User, backref='posts')

        Post.__table__.create()
        user = User(username='dude').save()
        post = Post(user=user, title='foo', body='bar').save()
        user.delete()
        self.assertEqual(User.query.count(), 0)
        self.assertEqual(Post.query.count(), 0)

    # Passes
    def test_backref_cascade_delete(self):
        Base = self.Base
        User = self.User
        class Post(Base):
            __tablename__ = 'posts'

            id = Column(Integer, primary_key=True)
            title = Column(String(255), nullable=False)
            body = Column(Text, nullable=False)
            user_id = Column(Integer, ForeignKey(User.id, ondelete='CASCADE'),
                             nullable=False)
            user = relationship(User,
                                backref=backref('posts', cascade='delete'))

        Post.__table__.create()
        user = User(username='dude').save()
        post = Post(user=user, title='foo', body='bar').save()
        user.delete()
        self.assertEqual(User.query.count(), 0)
        self.assertEqual(Post.query.count(), 0)