-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
139 lines (101 loc) · 3.21 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import enum
import os
try:
from psycopg2ct.compat import register
except ImportError:
pass
else:
register()
from pytest import fixture, mark, raises, yield_fixture
from sqlalchemy.engine import create_engine
from sqlalchemy.exc import StatementError
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import NullPool
from sqlalchemy.schema import Column
from sqlalchemy.types import Integer
from sqlalchemy_enum_list import (EnumListType, EnumListTypeException,
EnumSetType)
Base = declarative_base()
Session = sessionmaker()
try:
database_urls = os.environ['TEST_DATABASE_URLS'].split()
except KeyError:
database_urls = []
@fixture(scope='function', params=['sqlite://'] + database_urls)
def fx_engine(request):
url = request.param
engine = create_engine(url, poolclass=NullPool)
request.addfinalizer(engine.dispose)
return engine
@yield_fixture
def fx_connection(fx_engine):
connection = fx_engine.connect()
try:
transaction = connection.begin()
try:
metadata = Base.metadata
metadata.create_all(bind=connection)
yield connection
finally:
transaction.rollback()
finally:
connection.close()
@yield_fixture
def fx_session(fx_connection):
session = Session(bind=fx_connection)
try:
yield session
finally:
session.close()
class IntEnum(enum.Enum):
a = 1
b = 2
class StrEnum(enum.Enum):
a = 'a'
b = 'b'
class EnumTest(Base):
id = Column(Integer, primary_key=True)
int_column = Column(EnumListType(IntEnum, int))
str_column = Column(EnumListType(StrEnum, str))
str_column_set = Column(EnumSetType(StrEnum, str))
__tablename__ = 'enum_test'
@mark.parametrize(
'column_name, enum_cls',
[
('int_column', IntEnum),
('str_column', StrEnum),
]
)
def test_enum_list(fx_session, column_name, enum_cls):
entity = EnumTest(**{column_name: list(enum_cls)})
fx_session.add(entity)
fx_session.flush()
assert getattr(entity, column_name) == list(enum_cls)
def test_enum_list_init_exception(fx_session):
with raises(EnumListTypeException):
EnumListType(StrEnum, float, 'a')
with raises(EnumListTypeException):
EnumListType(StrEnum, float)
with raises(EnumListTypeException):
EnumListType(int, str)
def test_enum_list_bind_exception(fx_session):
fx_session.add(EnumTest(str_column='a'))
with raises(StatementError):
fx_session.flush()
def test_enum_set(fx_session):
entity = EnumTest(str_column_set={StrEnum.a, StrEnum.a, StrEnum.b})
fx_session.add(entity)
fx_session.flush()
assert entity.str_column_set == {StrEnum.a, StrEnum.b}
def test_enum_set_init_exception(fx_session):
with raises(EnumListTypeException):
EnumSetType(StrEnum, float, 'a')
with raises(EnumListTypeException):
EnumSetType(StrEnum, float)
with raises(EnumListTypeException):
EnumSetType(int, str)
def test_enum_set_bind_exception(fx_session):
fx_session.add(EnumTest(str_column_set={'a'}))
with raises(StatementError):
fx_session.flush()