1import pytest 2import sqlalchemy as sa 3 4from sqlalchemy_utils import get_primary_keys 5 6try: 7 from collections import OrderedDict 8except ImportError: 9 from ordereddict import OrderedDict 10 11 12@pytest.fixture 13def Building(Base): 14 class Building(Base): 15 __tablename__ = 'building' 16 id = sa.Column('_id', sa.Integer, primary_key=True) 17 name = sa.Column('_name', sa.Unicode(255)) 18 return Building 19 20 21class TestGetPrimaryKeys(object): 22 23 def test_table(self, Building): 24 assert get_primary_keys(Building.__table__) == OrderedDict({ 25 '_id': Building.__table__.c._id 26 }) 27 28 def test_declarative_class(self, Building): 29 assert get_primary_keys(Building) == OrderedDict({ 30 'id': Building.__table__.c._id 31 }) 32 33 def test_declarative_object(self, Building): 34 assert get_primary_keys(Building()) == OrderedDict({ 35 'id': Building.__table__.c._id 36 }) 37 38 def test_class_alias(self, Building): 39 alias = sa.orm.aliased(Building) 40 assert get_primary_keys(alias) == OrderedDict({ 41 'id': Building.__table__.c._id 42 }) 43 44 def test_table_alias(self, Building): 45 alias = sa.orm.aliased(Building.__table__) 46 assert get_primary_keys(alias) == OrderedDict({ 47 '_id': alias.c._id 48 }) 49