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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
|
from enum import Enum
from ipaddress import ip_network
from typing import Any, Dict, List
from uuid import UUID
import pytest
from sqlalchemy import ForeignKey, Text, __version__, orm, types
from sqlalchemy.dialects.mssql import JSON as MSSQL_JSON
from sqlalchemy.dialects.mysql import JSON as MYSQL_JSON
from sqlalchemy.dialects.postgresql import ARRAY, CIDR, HSTORE, INET, JSON, JSONB
from sqlalchemy.dialects.sqlite import JSON as SQLITE_JSON
from sqlalchemy.ext.mutable import MutableDict, MutableList
from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory
if __version__.startswith("1"):
pytest.importorskip("SQLAlchemy", "2")
class Base(orm.DeclarativeBase): ...
class Author(Base):
__tablename__ = "authors"
id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
books: orm.Mapped[List["Book"]] = orm.relationship(
"Book",
uselist=True,
back_populates="author",
)
class Book(Base):
__tablename__ = "books"
id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
author_id: orm.Mapped[int] = orm.mapped_column(ForeignKey(Author.id))
author: orm.Mapped[Author] = orm.relationship(
Author,
uselist=False,
back_populates="books",
)
def test_python_type_handling_v2() -> None:
class Base(orm.DeclarativeBase): ...
class Animal(str, Enum):
DOG = "Dog"
CAT = "Cat"
class Model(Base):
__tablename__ = "model"
id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
str_type: orm.Mapped[str]
enum_type: orm.Mapped[Animal]
str_array_type: orm.Mapped[List[str]] = orm.mapped_column(
type_=types.ARRAY(types.String),
)
class ModelFactory(SQLAlchemyFactory[Model]):
__model__ = Model
instance = ModelFactory.build()
assert isinstance(instance.id, int)
assert isinstance(instance.str_type, str)
assert isinstance(instance.enum_type, Animal)
assert isinstance(instance.str_array_type, list)
assert isinstance(instance.str_array_type[0], str)
def test_sqla_dialect_types() -> None:
class Base(orm.DeclarativeBase): ...
class SqlaModel(Base):
__tablename__ = "sql_models"
id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
uuid_type: orm.Mapped[UUID] = orm.mapped_column(type_=types.UUID)
nested_array_inet: orm.Mapped[List[str]] = orm.mapped_column(type_=ARRAY(INET, dimensions=1))
nested_array_cidr: orm.Mapped[List[str]] = orm.mapped_column(type_=ARRAY(CIDR, dimensions=1))
hstore_type: orm.Mapped[Dict] = orm.mapped_column(type_=HSTORE)
mut_nested_array_inet: orm.Mapped[List[str]] = orm.mapped_column(
type_=MutableList.as_mutable(ARRAY(INET, dimensions=1))
)
pg_json_type: orm.Mapped[Dict] = orm.mapped_column(type_=JSON)
pg_jsonb_type: orm.Mapped[Dict] = orm.mapped_column(type_=JSONB)
common_json_type: orm.Mapped[Dict] = orm.mapped_column(type_=types.JSON)
mysql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MYSQL_JSON)
sqlite_json: orm.Mapped[Dict] = orm.mapped_column(type_=SQLITE_JSON)
mssql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MSSQL_JSON)
multable_pg_json_type: orm.Mapped[Dict] = orm.mapped_column(
type_=MutableDict.as_mutable(JSON(astext_type=Text()))
)
multable_pg_jsonb_type: orm.Mapped[Dict] = orm.mapped_column(
type_=MutableDict.as_mutable(JSONB(astext_type=Text()))
)
multable_common_json_type: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(types.JSON()))
multable_mysql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(MYSQL_JSON()))
multable_sqlite_json: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(SQLITE_JSON()))
multable_mssql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(MSSQL_JSON()))
class ModelFactory(SQLAlchemyFactory[SqlaModel]):
__model__ = SqlaModel
instance = ModelFactory.build()
assert isinstance(instance.nested_array_inet[0], str)
assert ip_network(instance.nested_array_inet[0])
assert isinstance(instance.nested_array_cidr[0], str)
assert ip_network(instance.nested_array_cidr[0])
assert isinstance(instance.hstore_type, dict)
assert isinstance(instance.uuid_type, UUID)
assert isinstance(instance.mut_nested_array_inet[0], str)
assert ip_network(instance.mut_nested_array_inet[0])
assert isinstance(instance.pg_json_type, dict)
for value in instance.pg_json_type.values():
assert isinstance(value, (str, int, bool, float))
assert isinstance(instance.pg_jsonb_type, dict)
for value in instance.pg_jsonb_type.values():
assert isinstance(value, (str, int, bool, float))
assert isinstance(instance.common_json_type, dict)
for value in instance.common_json_type.values():
assert isinstance(value, (str, int, bool, float))
assert isinstance(instance.mysql_json, dict)
for value in instance.mysql_json.values():
assert isinstance(value, (str, int, bool, float))
assert isinstance(instance.sqlite_json, dict)
for value in instance.sqlite_json.values():
assert isinstance(value, (str, int, bool, float))
assert isinstance(instance.mssql_json, dict)
for value in instance.mssql_json.values():
assert isinstance(value, (str, int, bool, float))
assert isinstance(instance.multable_pg_json_type, dict)
assert isinstance(instance.multable_pg_jsonb_type, dict)
assert isinstance(instance.multable_common_json_type, dict)
assert isinstance(instance.multable_mysql_json, dict)
assert isinstance(instance.multable_sqlite_json, dict)
assert isinstance(instance.multable_mssql_json, dict)
@pytest.mark.parametrize(
"type_",
tuple(SQLAlchemyFactory.get_sqlalchemy_types().keys()),
)
def test_sqlalchemy_type_handlers_v2(type_: types.TypeEngine) -> None:
class Base(orm.DeclarativeBase): ...
class Model(Base):
__tablename__ = "model_with_overridden_type"
id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
overridden: orm.Mapped[Any] = orm.mapped_column(type_=type_)
class ModelFactory(SQLAlchemyFactory[Model]):
__model__ = Model
instance = ModelFactory.build()
assert instance.overridden is not None
|
