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
|
import importlib
import pytest
from sqlalchemy import create_mock_engine
from sqlalchemy.sql.type_api import TypeEngine
from sqlmodel import SQLModel
from . import test_enums_models
from .conftest import needs_pydanticv1, needs_pydanticv2
"""
Tests related to Enums
Associated issues:
* https://github.com/tiangolo/sqlmodel/issues/96
* https://github.com/tiangolo/sqlmodel/issues/164
"""
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(clear_sqlmodel, capsys: pytest.CaptureFixture[str]):
assert test_enums_models, "Ensure the models are imported and registered"
importlib.reload(test_enums_models)
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(clear_sqlmodel, capsys: pytest.CaptureFixture[str]):
assert test_enums_models, "Ensure the models are imported and registered"
importlib.reload(test_enums_models)
SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False)
captured = capsys.readouterr()
assert "enum_field VARCHAR(1) NOT NULL" in captured.out, captured
assert "CREATE TYPE" not in captured.out
@needs_pydanticv1
def test_json_schema_flat_model_pydantic_v1():
assert test_enums_models.FlatModel.schema() == {
"title": "FlatModel",
"type": "object",
"properties": {
"id": {"title": "Id", "type": "string", "format": "uuid"},
"enum_field": {"$ref": "#/definitions/MyEnum1"},
},
"required": ["id", "enum_field"],
"definitions": {
"MyEnum1": {
"title": "MyEnum1",
"description": "An enumeration.",
"enum": ["A", "B"],
"type": "string",
}
},
}
@needs_pydanticv1
def test_json_schema_inherit_model_pydantic_v1():
assert test_enums_models.InheritModel.schema() == {
"title": "InheritModel",
"type": "object",
"properties": {
"id": {"title": "Id", "type": "string", "format": "uuid"},
"enum_field": {"$ref": "#/definitions/MyEnum2"},
},
"required": ["id", "enum_field"],
"definitions": {
"MyEnum2": {
"title": "MyEnum2",
"description": "An enumeration.",
"enum": ["C", "D"],
"type": "string",
}
},
}
@needs_pydanticv2
def test_json_schema_flat_model_pydantic_v2():
assert test_enums_models.FlatModel.model_json_schema() == {
"title": "FlatModel",
"type": "object",
"properties": {
"id": {"title": "Id", "type": "string", "format": "uuid"},
"enum_field": {"$ref": "#/$defs/MyEnum1"},
},
"required": ["id", "enum_field"],
"$defs": {
"MyEnum1": {"enum": ["A", "B"], "title": "MyEnum1", "type": "string"}
},
}
@needs_pydanticv2
def test_json_schema_inherit_model_pydantic_v2():
assert test_enums_models.InheritModel.model_json_schema() == {
"title": "InheritModel",
"type": "object",
"properties": {
"id": {"title": "Id", "type": "string", "format": "uuid"},
"enum_field": {"$ref": "#/$defs/MyEnum2"},
},
"required": ["id", "enum_field"],
"$defs": {
"MyEnum2": {"enum": ["C", "D"], "title": "MyEnum2", "type": "string"}
},
}
|
