From 0558934486923ceffc1920851203ebcaa6dff9e8 Mon Sep 17 00:00:00 2001 From: rusty Date: Sat, 17 Feb 2024 19:45:17 +0400 Subject: [PATCH] Add more tests that use CompareAndRunTestCase --- tests/base/run_migration_test_abc.py | 4 + tests/test_alter_column/test_text_column.py | 86 +++++++++++---------- tests/test_enum_creation/test_add_column.py | 27 +++++++ 3 files changed, 77 insertions(+), 40 deletions(-) diff --git a/tests/base/run_migration_test_abc.py b/tests/base/run_migration_test_abc.py index fb2b7cb..93430c2 100644 --- a/tests/base/run_migration_test_abc.py +++ b/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, diff --git a/tests/test_alter_column/test_text_column.py b/tests/test_alter_column/test_text_column.py index d213963..404f230 100644 --- a/tests/test_alter_column/test_text_column.py +++ b/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 ### + """ diff --git a/tests/test_enum_creation/test_add_column.py b/tests/test_enum_creation/test_add_column.py index 393fa6d..6716060 100644 --- a/tests/test_enum_creation/test_add_column.py +++ b/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()