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