polymorphic images

images.py

import pytest

from sqlalchemy import Column, Integer, ForeignKey, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship


class Base(object):
    pass


DeclarativeBase = declarative_base(cls=Base)


class Property(DeclarativeBase):
    __tablename__ = "properties"

    id = Column(Integer, primary_key=True)
    name = Column(String)


class UnitType(DeclarativeBase):
    __tablename__ = "unit_types"

    id = Column(Integer, primary_key=True)
    name = Column(String)


class Image(DeclarativeBase):
    __tablename__ = "images"

    id = Column(Integer, primary_key=True)

    discriminator = Column('type', String(50))
    __mapper_args__ = {'polymorphic_on': discriminator}


class PropertyImageBase(Image):
    property_id = Column(
        Integer,
        ForeignKey("properties.id", name="fk_property_image_base_properties")
    )
    property = relationship(Property, backref="images")


class CityImageBase(Image):
    pass


class CityImage(CityImageBase):
    pass


class CityPointsOfInterestImage(CityImageBase):
    pass


class PropertyImage(PropertyImageBase):
    __mapper_args__ = {'polymorphic_identity': 'property_image'}
    property = relationship(Property, backref="property_images")


class UnitTypeImage(PropertyImageBase):
    unit_type_id = Column(
        Integer,
        ForeignKey("unit_types.id", name="fk_unit_type_image_unit_type")
    )

    __mapper_args__ = {'polymorphic_identity': 'unit_type_image'}
    property = relationship(Property, backref="unit_type_images")
    unit_type = relationship(UnitType, backref="images")


class PropertyPointsOfInterestImage(PropertyImageBase):
    __mapper_args__ = {
        'polymorphic_identity': 'property_points_of_interest_image'
    }
    property = relationship(Property, backref="points_of_interest_images")


class PropertyQuickFactsImage(PropertyImageBase):
    __mapper_args__ = {
        'polymorphic_identity': 'property_quick_facts_image'
    }
    property = relationship(Property, backref="quick_facts_images")


@pytest.fixture(scope='session')
def model_base():
    return DeclarativeBase


def test_property_images(db_session):

    property_ = Property(name="Student Lounge")
    unit_type = UnitType(name="Sweet suit")
    property_image = PropertyImage(property=property_)

    property_points_of_interest_image = PropertyPointsOfInterestImage(
        property=property_
    )

    property_quick_facts_image = PropertyQuickFactsImage(
        property=property_
    )

    unit_type_image = UnitTypeImage(
        property=property_,
        unit_type=unit_type
    )

    db_session.add_all([
        property_image,
        property_points_of_interest_image,
        property_quick_facts_image,
        unit_type_image
    ])
    db_session.commit()

    assert len(property_.images) == 4
    assert len(property_.property_images) == 1
    assert len(property_.unit_type_images) == 1
    assert len(unit_type.images) == 1
    assert len(property_.points_of_interest_images) == 1
    assert len(property_.quick_facts_images) == 1