Skip to content

Commit

Permalink
✅ Refactor test_enums to make them independent of previous imports (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
tiangolo authored Aug 31, 2024
1 parent 9acd934 commit e4f3ec7
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 34 deletions.
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

0 comments on commit e4f3ec7

Please sign in to comment.