Merge pull request #72 from Pogchamp-company/feature/-/tests_base_class

Create tests base class

tests/base/render_and_run.py

@@ -20,6 +20,7 @@ def compare_and_run(
     *,
     expected_upgrade: str,
     expected_downgrade: str,
+    disable_running: bool = False,
 ):
     """Compares generated migration script is equal to expected_upgrade and expected_downgrade, then runs it"""
     migration_context = create_migration_context(connection, target_schema)
@@ -39,6 +40,9 @@ def compare_and_run(
     assert upgrade_code == expected_upgrade, f"Got:\n{upgrade_code!r}\nExpected:\n{expected_upgrade!r}"
     assert downgrade_code == expected_downgrade, f"Got:\n{downgrade_code!r}\nExpected:\n{expected_downgrade!r}"
 
+    if disable_running:
+        return
+
     exec(
         upgrade_code,
         {  # todo Use imports from template_args

tests/base/run_migration_test_abc.py

@@ -9,12 +9,21 @@
 
 
 class CompareAndRunTestCase(ABC):
+    """
+    Base class for all tests that expect specific alembic generated code
+    """
+
+    disable_running = False
+
     @abstractmethod
     def get_database_schema(self) -> MetaData: ...
 
     @abstractmethod
     def get_target_schema(self) -> MetaData: ...
 
+    def insert_migration_data(self, connection: "Connection", database_schema: MetaData) -> None:
+        pass
+
     @abstractmethod
     def get_expected_upgrade(self) -> str: ...
 
@@ -26,10 +35,12 @@ def test_run(self, connection: "Connection"):
         target_schema = self.get_target_schema()
 
         database_schema.create_all(connection)
+        self.insert_migration_data(connection, database_schema)
 
         compare_and_run(
             connection,
             target_schema,
             expected_upgrade=self.get_expected_upgrade(),
             expected_downgrade=self.get_expected_downgrade(),
+            disable_running=self.disable_running,
         )

tests/test_alter_column/test_text_column.py

@@ -4,56 +4,61 @@
 from sqlalchemy import MetaData, Table, Column, TEXT, insert
 from sqlalchemy.dialects import postgresql
 
+from tests.base.run_migration_test_abc import CompareAndRunTestCase
+
 if TYPE_CHECKING:
     from sqlalchemy import Connection
 
-from tests.base.render_and_run import compare_and_run
-
 
 class NewEnum(Enum):
     A = "a"
     B = "b"
     C = "c"
 
 
-def test_text_column(connection: "Connection"):
-    database_schema = MetaData()
-    a_table = Table("a", database_schema, Column("value", TEXT))
-    database_schema.create_all(connection)
-    connection.execute(
-        insert(a_table).values(
-            [
-                {"value": NewEnum.A.name},
-                {"value": NewEnum.B.name},
-                {"value": NewEnum.B.name},
-                {"value": NewEnum.C.name},
-            ]
+class TestTextColumn(CompareAndRunTestCase):
+    def get_database_schema(self) -> MetaData:
+        database_schema = MetaData()
+        Table("a", database_schema, Column("value", TEXT))
+        return database_schema
+
+    def get_target_schema(self) -> MetaData:
+        target_schema = MetaData()
+        Table("a", target_schema, Column("value", postgresql.ENUM(NewEnum)))
+        return target_schema
+
+    def insert_migration_data(self, connection: "Connection", database_schema: MetaData) -> None:
+        a_table = database_schema.tables["a"]
+        connection.execute(
+            insert(a_table).values(
+                [
+                    {"value": NewEnum.A.name},
+                    {"value": NewEnum.B.name},
+                    {"value": NewEnum.B.name},
+                    {"value": NewEnum.C.name},
+                ]
+            )
         )
-    )
-
-    target_schema = MetaData()
-    Table("a", target_schema, Column("value", postgresql.ENUM(NewEnum)))
-
-    compare_and_run(
-        connection,
-        target_schema,
-        expected_upgrade=f"""
-        # ### commands auto generated by Alembic - please adjust! ###
-        sa.Enum('A', 'B', 'C', name='newenum').create(op.get_bind())
-        op.alter_column('a', 'value',
-                   existing_type=sa.TEXT(),
-                   type_=postgresql.ENUM('A', 'B', 'C', name='newenum'),
-                   existing_nullable=True,
-                   postgresql_using='value::newenum')
-        # ### end Alembic commands ###
-    """,
-        expected_downgrade=f"""
-        # ### commands auto generated by Alembic - please adjust! ###
-        op.alter_column('a', 'value',
-                   existing_type=postgresql.ENUM('A', 'B', 'C', name='newenum'),
-                   type_=sa.TEXT(),
-                   existing_nullable=True)
-        sa.Enum('A', 'B', 'C', name='newenum').drop(op.get_bind())
-        # ### end Alembic commands ###    
-    """,
-    )
+
+    def get_expected_upgrade(self) -> str:
+        return """
+            # ### commands auto generated by Alembic - please adjust! ###
+            sa.Enum('A', 'B', 'C', name='newenum').create(op.get_bind())
+            op.alter_column('a', 'value',
+                       existing_type=sa.TEXT(),
+                       type_=postgresql.ENUM('A', 'B', 'C', name='newenum'),
+                       existing_nullable=True,
+                       postgresql_using='value::newenum')
+            # ### end Alembic commands ###
+        """
+
+    def get_expected_downgrade(self) -> str:
+        return """
+            # ### commands auto generated by Alembic - please adjust! ###
+            op.alter_column('a', 'value',
+                       existing_type=postgresql.ENUM('A', 'B', 'C', name='newenum'),
+                       type_=sa.TEXT(),
+                       existing_nullable=True)
+            sa.Enum('A', 'B', 'C', name='newenum').drop(op.get_bind())
+            # ### end Alembic commands ###    
+        """

tests/test_enum_creation/test_add_column.py

@@ -6,7 +6,7 @@
 from alembic.operations import ops
 
 from alembic_postgresql_enum.operations import CreateEnumOp
-from tests.base.render_and_run import compare_and_run
+from tests.base.run_migration_test_abc import CompareAndRunTestCase
 from tests.schemas import (
     get_schema_with_enum_variants,
     USER_TABLE_NAME,
@@ -25,31 +25,58 @@
 from sqlalchemy import MetaData, Table, Column, Integer
 
 
-def test_create_enum_before_add_column(connection: "Connection"):
+class TestCreateEnumBeforeAddColumn(CompareAndRunTestCase):
+    new_enum_variants = ["active", "passive"]
+
+    def get_database_schema(self) -> MetaData:
+        return get_schema_without_enum()
+
+    def get_target_schema(self) -> MetaData:
+        return get_schema_with_enum_variants(self.new_enum_variants)
+
+    def get_expected_upgrade(self) -> str:
+        return f"""
+        # ### commands auto generated by Alembic - please adjust! ###
+        sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').create(op.get_bind())
+        op.add_column('{USER_TABLE_NAME}', sa.Column('{USER_STATUS_COLUMN_NAME}', postgresql.ENUM({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', create_type=False), nullable=True))
+        # ### end Alembic commands ###
+        """
+
+    def get_expected_downgrade(self) -> str:
+        return f"""
+        # ### commands auto generated by Alembic - please adjust! ###
+        op.drop_column('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')
+        sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind())
+        # ### end Alembic commands ###
+        """
+
+
+class TestCreateEnumBeforeAddColumn(CompareAndRunTestCase):
     """Check that library correctly creates enum before its use inside add_column"""
-    database_schema = get_schema_without_enum()
-    database_schema.create_all(connection)
 
     new_enum_variants = ["active", "passive"]
 
-    target_schema = get_schema_with_enum_variants(new_enum_variants)
+    def get_database_schema(self) -> MetaData:
+        return get_schema_without_enum()
+
+    def get_target_schema(self) -> MetaData:
+        return get_schema_with_enum_variants(self.new_enum_variants)
 
-    compare_and_run(
-        connection,
-        target_schema,
-        expected_upgrade=f"""
+    def get_expected_upgrade(self) -> str:
+        return f"""
         # ### commands auto generated by Alembic - please adjust! ###
-        sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').create(op.get_bind())
-        op.add_column('{USER_TABLE_NAME}', sa.Column('{USER_STATUS_COLUMN_NAME}', postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', create_type=False), nullable=True))
+        sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').create(op.get_bind())
+        op.add_column('{USER_TABLE_NAME}', sa.Column('{USER_STATUS_COLUMN_NAME}', postgresql.ENUM({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', create_type=False), nullable=True))
         # ### end Alembic commands ###
-    """,
-        expected_downgrade=f"""
+        """
+
+    def get_expected_downgrade(self) -> str:
+        return f"""
         # ### commands auto generated by Alembic - please adjust! ###
         op.drop_column('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')
-        sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind())
+        sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind())
         # ### end Alembic commands ###
-    """,
-    )
+        """
 
 
 def test_create_enum_diff_tuple(connection: "Connection"):
@@ -110,63 +137,67 @@ def test_create_enum_diff_tuple_with_array(connection: "Connection"):
     assert add_column_tuple[0] == "add_column"
 
 
-def test_with_non_native_enum(connection: "Connection"):
+class TestWithNonNativeEnum(CompareAndRunTestCase):
     """Check that library ignores sa.Enum that are not native"""
-    database_schema = get_schema_without_enum()
-    database_schema.create_all(connection)
 
     new_enum_variants = ["active", "passive"]
 
-    target_schema = MetaData()
+    def get_database_schema(self) -> MetaData:
+        return get_schema_without_enum()
 
-    Table(
-        USER_TABLE_NAME,
-        target_schema,
-        Column("id", Integer, primary_key=True),
-        Column(
-            USER_STATUS_COLUMN_NAME,
-            sqlalchemy.Enum(*new_enum_variants, name=USER_STATUS_ENUM_NAME, native_enum=False),
-        ),
-    )
+    def get_target_schema(self) -> MetaData:
+        target_schema = MetaData()
+
+        Table(
+            USER_TABLE_NAME,
+            target_schema,
+            Column("id", Integer, primary_key=True),
+            Column(
+                USER_STATUS_COLUMN_NAME,
+                sqlalchemy.Enum(*self.new_enum_variants, name=USER_STATUS_ENUM_NAME, native_enum=False),
+            ),
+        )
 
-    compare_and_run(
-        connection,
-        target_schema,
-        expected_upgrade=f"""
+        return target_schema
+
+    def get_expected_upgrade(self) -> str:
+        return f"""
         # ### commands auto generated by Alembic - please adjust! ###
-        op.add_column('{USER_TABLE_NAME}', sa.Column('{USER_STATUS_COLUMN_NAME}', sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', native_enum=False), nullable=True))
+        op.add_column('{USER_TABLE_NAME}', sa.Column('{USER_STATUS_COLUMN_NAME}', sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', native_enum=False), nullable=True))
         # ### end Alembic commands ###
-    """,
-        expected_downgrade=f"""
+        """
+
+    def get_expected_downgrade(self) -> str:
+        return f"""
         # ### commands auto generated by Alembic - please adjust! ###
         op.drop_column('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')
         # ### end Alembic commands ###
-    """,
-    )
+        """
 
 
-def test_create_enum_before_add_column_metadata_list(connection: "Connection"):
+class TestCreateEnumBeforeAddColumnMetadataList(CompareAndRunTestCase):
     """Check that library correctly creates enum before its use inside add_column when metadata is in list"""
-    database_schema = get_schema_without_enum()
-    database_schema.create_all(connection)
 
     new_enum_variants = ["active", "passive"]
 
-    target_schema = get_schema_with_enum_variants(new_enum_variants)
+    def get_database_schema(self) -> MetaData:
+        return get_schema_without_enum()
 
-    compare_and_run(
-        connection,
-        [target_schema],
-        expected_upgrade=f"""
+    def get_target_schema(self) -> MetaData:
+        return get_schema_with_enum_variants(self.new_enum_variants)
+
+    def get_expected_upgrade(self) -> str:
+        return f"""
         # ### commands auto generated by Alembic - please adjust! ###
-        sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').create(op.get_bind())
-        op.add_column('{USER_TABLE_NAME}', sa.Column('{USER_STATUS_COLUMN_NAME}', postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', create_type=False), nullable=True))
+        sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').create(op.get_bind())
+        op.add_column('{USER_TABLE_NAME}', sa.Column('{USER_STATUS_COLUMN_NAME}', postgresql.ENUM({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', create_type=False), nullable=True))
         # ### end Alembic commands ###
-    """,
-        expected_downgrade=f"""
+        """
+
+    def get_expected_downgrade(self) -> str:
+        return f"""
         # ### commands auto generated by Alembic - please adjust! ###
         op.drop_column('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')
-        sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind())
+        sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind())
         # ### end Alembic commands ###
-    """,
-    )
+        """