Reduced code duplication in tests using parametrization

Author: YuriiMotov

tests/test_field_sa_column.py

@@ -1,78 +1,86 @@
-from typing import Optional
+from typing import Optional, Union
 
 import pytest
 from sqlalchemy import Column, Integer, String
+from sqlalchemy.orm import mapped_column
 from sqlmodel import Field, SQLModel
 
 
-def test_sa_column_takes_precedence(clear_sqlmodel) -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_takes_precedence(clear_sqlmodel, column_class) -> None:
     class Item(SQLModel, table=True):
         id: Optional[int] = Field(
             default=None,
-            sa_column=Column(String, primary_key=True, nullable=False),
+            sa_column=column_class(String, primary_key=True, nullable=False),
         )
 
     # It would have been nullable with no sa_column
     assert Item.id.nullable is False  # type: ignore
     assert isinstance(Item.id.type, String)  # type: ignore
 
 
-def test_sa_column_no_sa_args() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_sa_args(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
                 sa_column_args=[Integer],
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_sa_kargs() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_sa_kargs(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
                 sa_column_kwargs={"primary_key": True},
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_type() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_type(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
                 sa_type=Integer,
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_primary_key() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_primary_key(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
                 primary_key=True,
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_nullable() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_nullable(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
                 nullable=True,
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_foreign_key() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_foreign_key(clear_sqlmodel, column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Team(SQLModel, table=True):
@@ -84,38 +92,41 @@ class Hero(SQLModel, table=True):
             team_id: Optional[int] = Field(
                 default=None,
                 foreign_key="team.id",
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_unique() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_unique(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
                 unique=True,
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_index() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_index(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
                 index=True,
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
             )
 
 
-def test_sa_column_no_ondelete() -> None:
+@pytest.mark.parametrize("column_class", [Column, mapped_column])
+def test_sa_column_no_ondelete(column_class) -> None:
     with pytest.raises(RuntimeError):
 
         class Item(SQLModel, table=True):
             id: Optional[int] = Field(
                 default=None,
-                sa_column=Column(Integer, primary_key=True),
+                sa_column=column_class(Integer, primary_key=True),
                 ondelete="CASCADE",
             )