Add more tests that use CompareAndRunTestCase

tests/base/run_migration_test_abc.py

@@ -15,6 +15,9 @@ def get_database_schema(self) -> MetaData: ...
     @abstractmethod
     def get_target_schema(self) -> MetaData: ...
 
+    def insert_migration_data(self, connection: "Connection"):
+        pass
+
     @abstractmethod
     def get_expected_upgrade(self) -> str: ...
 
@@ -26,6 +29,7 @@ def test_run(self, connection: "Connection"):
         target_schema = self.get_target_schema()
 
         database_schema.create_all(connection)
+        self.insert_migration_data(connection)
 
         compare_and_run(
             connection,

tests/test_alter_column/test_text_column.py

@@ -4,6 +4,8 @@
 from sqlalchemy import MetaData, Table, Column, TEXT, insert
 from sqlalchemy.dialects import postgresql
 
+from base.run_migration_test_abc import CompareAndRunTestCase
+
 if TYPE_CHECKING:
     from sqlalchemy import Connection
 
@@ -16,44 +18,48 @@ class NewEnum(Enum):
     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()
+        self.a_table = 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"):
+        connection.execute(
+            insert(self.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,6 +6,7 @@
 from alembic.operations import ops
 
 from alembic_postgresql_enum.operations import CreateEnumOp
+from base.run_migration_test_abc import CompareAndRunTestCase
 from tests.base.render_and_run import compare_and_run
 from tests.schemas import (
     get_schema_with_enum_variants,
@@ -25,6 +26,32 @@
 from sqlalchemy import MetaData, Table, Column, Integer
 
 
+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 ###
+        """
+
+
 def test_create_enum_before_add_column(connection: "Connection"):
     """Check that library correctly creates enum before its use inside add_column"""
     database_schema = get_schema_without_enum()