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
|
import enum
import uuid
from sqlalchemy import create_mock_engine
from sqlalchemy.sql.type_api import TypeEngine
from sqlmodel import Field, SQLModel
"""
Tests related to Enums
Associated issues:
* https://github.com/tiangolo/sqlmodel/issues/96
* https://github.com/tiangolo/sqlmodel/issues/164
"""
class MyEnum1(enum.Enum):
A = "A"
B = "B"
class MyEnum2(enum.Enum):
C = "C"
D = "D"
class BaseModel(SQLModel):
id: uuid.UUID = Field(primary_key=True)
enum_field: MyEnum2
class FlatModel(SQLModel, table=True):
id: uuid.UUID = Field(primary_key=True)
enum_field: MyEnum1
class InheritModel(BaseModel, table=True):
pass
def pg_dump(sql: TypeEngine, *args, **kwargs):
dialect = sql.compile(dialect=postgres_engine.dialect)
sql_str = str(dialect).rstrip()
if sql_str:
print(sql_str + ";")
def sqlite_dump(sql: TypeEngine, *args, **kwargs):
dialect = sql.compile(dialect=sqlite_engine.dialect)
sql_str = str(dialect).rstrip()
if sql_str:
print(sql_str + ";")
postgres_engine = create_mock_engine("postgresql://", pg_dump)
sqlite_engine = create_mock_engine("sqlite://", sqlite_dump)
def test_postgres_ddl_sql(capsys):
SQLModel.metadata.create_all(bind=postgres_engine, checkfirst=False)
captured = capsys.readouterr()
assert "CREATE TYPE myenum1 AS ENUM ('A', 'B');" in captured.out
assert "CREATE TYPE myenum2 AS ENUM ('C', 'D');" in captured.out
def test_sqlite_ddl_sql(capsys):
SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False)
captured = capsys.readouterr()
assert "enum_field VARCHAR(1) NOT NULL" in captured.out
assert "CREATE TYPE" not in captured.out
|
