Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 15 additions & 34 deletions tests/test_enums.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import enum
import uuid
import importlib

import pytest
from sqlalchemy import create_mock_engine
from sqlalchemy.sql.type_api import TypeEngine
from sqlmodel import Field, SQLModel
from sqlmodel import SQLModel

from . import test_enums_models
from .conftest import needs_pydanticv1, needs_pydanticv2

"""
Expand All @@ -16,30 +17,6 @@
"""


class MyEnum1(str, enum.Enum):
A = "A"
B = "B"


class MyEnum2(str, 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()
Expand All @@ -58,25 +35,29 @@ def sqlite_dump(sql: TypeEngine, *args, **kwargs):
sqlite_engine = create_mock_engine("sqlite://", sqlite_dump)


def test_postgres_ddl_sql(capsys):
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(capsys):
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
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 FlatModel.schema() == {
assert test_enums_models.FlatModel.schema() == {
"title": "FlatModel",
"type": "object",
"properties": {
Expand All @@ -97,7 +78,7 @@ def test_json_schema_flat_model_pydantic_v1():

@needs_pydanticv1
def test_json_schema_inherit_model_pydantic_v1():
assert InheritModel.schema() == {
assert test_enums_models.InheritModel.schema() == {
"title": "InheritModel",
"type": "object",
"properties": {
Expand All @@ -118,7 +99,7 @@ def test_json_schema_inherit_model_pydantic_v1():

@needs_pydanticv2
def test_json_schema_flat_model_pydantic_v2():
assert FlatModel.model_json_schema() == {
assert test_enums_models.FlatModel.model_json_schema() == {
"title": "FlatModel",
"type": "object",
"properties": {
Expand All @@ -134,7 +115,7 @@ def test_json_schema_flat_model_pydantic_v2():

@needs_pydanticv2
def test_json_schema_inherit_model_pydantic_v2():
assert InheritModel.model_json_schema() == {
assert test_enums_models.InheritModel.model_json_schema() == {
"title": "InheritModel",
"type": "object",
"properties": {
Expand Down
28 changes: 28 additions & 0 deletions tests/test_enums_models.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
import enum
import uuid

from sqlmodel import Field, SQLModel


class MyEnum1(str, enum.Enum):
A = "A"
B = "B"


class MyEnum2(str, 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
Loading