✅ Refactor test_enums to make them independent of previous imports (#1095)

Author: tiangolo

tests/test_enums.py

@@ -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
 
 """
@@ -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()
@@ -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": {
@@ -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": {
@@ -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": {
@@ -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": {

tests/test_enums_models.py

@@ -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