1import pytest 2import sqlalchemy as sa 3 4from sqlalchemy_utils import get_columns 5 6 7@pytest.fixture 8def Building(Base): 9 class Building(Base): 10 __tablename__ = 'building' 11 id = sa.Column('_id', sa.Integer, primary_key=True) 12 name = sa.Column('_name', sa.Unicode(255)) 13 return Building 14 15 16@pytest.fixture 17def columns(): 18 return ['_id', '_name'] 19 20 21class TestGetColumns(object): 22 def test_table(self, Building): 23 assert isinstance( 24 get_columns(Building.__table__), 25 sa.sql.base.ImmutableColumnCollection 26 ) 27 28 def test_instrumented_attribute(self, Building): 29 assert get_columns(Building.id) == [Building.__table__.c._id] 30 31 def test_column_property(self, Building): 32 assert get_columns(Building.id.property) == [ 33 Building.__table__.c._id 34 ] 35 36 def test_column(self, Building): 37 assert get_columns(Building.__table__.c._id) == [ 38 Building.__table__.c._id 39 ] 40 41 def test_declarative_class(self, Building, columns): 42 assert [c.name for c in get_columns(Building).values()] == columns 43 44 def test_declarative_object(self, Building, columns): 45 assert [c.name for c in get_columns(Building()).values()] == columns 46 47 def test_mapper(self, Building, columns): 48 assert [ 49 c.name for c in get_columns(Building.__mapper__).values() 50 ] == columns 51 52 def test_class_alias(self, Building, columns): 53 assert [ 54 c.name for c in get_columns(sa.orm.aliased(Building)).values() 55 ] == columns 56 57 def test_table_alias(self, Building, columns): 58 alias = sa.orm.aliased(Building.__table__) 59 assert [c.name for c in get_columns(alias).values()] == columns 60