diff --git a/tests/base/render_and_run.py b/tests/base/render_and_run.py index ff39f60..a4e5fe6 100644 --- a/tests/base/render_and_run.py +++ b/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 diff --git a/tests/base/run_migration_test_abc.py b/tests/base/run_migration_test_abc.py index b310b26..eb6069a 100644 --- a/tests/base/run_migration_test_abc.py +++ b/tests/base/run_migration_test_abc.py @@ -9,24 +9,26 @@ class CompareAndRunTestCase(ABC): + """ + Base class for all tests that expect specific alembic generated code + """ + + disable_running = False + @abstractmethod - def get_database_schema(self) -> MetaData: - ... + def get_database_schema(self) -> MetaData: ... @abstractmethod - def get_target_schema(self) -> MetaData: - ... + 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: - ... + def get_expected_upgrade(self) -> str: ... @abstractmethod - def get_expected_downgrade(self) -> str: - ... + def get_expected_downgrade(self) -> str: ... def test_run(self, connection: "Connection"): database_schema = self.get_database_schema() @@ -40,4 +42,5 @@ def test_run(self, connection: "Connection"): target_schema, expected_upgrade=self.get_expected_upgrade(), expected_downgrade=self.get_expected_downgrade(), + disable_running=self.disable_running, ) diff --git a/tests/test_enum_creation/test_add_column.py b/tests/test_enum_creation/test_add_column.py index fd0649e..a97aa54 100644 --- a/tests/test_enum_creation/test_add_column.py +++ b/tests/test_enum_creation/test_add_column.py @@ -7,7 +7,6 @@ from alembic_postgresql_enum.operations import CreateEnumOp from tests.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, USER_TABLE_NAME, @@ -52,31 +51,32 @@ def get_expected_downgrade(self) -> str: """ -def test_create_enum_before_add_column(connection: "Connection"): +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"): @@ -137,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), + ), + ) + + return target_schema - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + 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 ### - """, - ) + """ diff --git a/tests/test_enum_creation/test_create_array.py b/tests/test_enum_creation/test_create_array.py index bc67845..c8d440d 100644 --- a/tests/test_enum_creation/test_create_array.py +++ b/tests/test_enum_creation/test_create_array.py @@ -1,8 +1,9 @@ -from typing import TYPE_CHECKING, List +from typing import List import sqlalchemy +from sqlalchemy import Table, Column, Integer, MetaData -from tests.base.render_and_run import compare_and_run +from tests.base.run_migration_test_abc import CompareAndRunTestCase from tests.schemas import ( get_car_schema_without_enum, CAR_TABLE_NAME, @@ -10,10 +11,6 @@ CAR_COLORS_ENUM_NAME, ) -if TYPE_CHECKING: - from sqlalchemy import Connection -from sqlalchemy import Table, Column, Integer, MetaData - def get_schema_with_enum_in_sqlalchemy_array_variants(variants: List[str]) -> MetaData: schema = MetaData() @@ -31,28 +28,29 @@ def get_schema_with_enum_in_sqlalchemy_array_variants(variants: List[str]) -> Me return schema -def test_create_enum_on_create_table_with_array(connection: "Connection"): +class TestCreateEnumOnCreateTableWithArray(CompareAndRunTestCase): """Check that library correctly creates enum before its use inside create_table. Enum is used in ARRAY""" - database_schema = get_car_schema_without_enum() - database_schema.create_all(connection) new_enum_variants = ["black", "white", "red", "green", "blue", "other"] - target_schema = get_schema_with_enum_in_sqlalchemy_array_variants(new_enum_variants) + def get_database_schema(self) -> MetaData: + return get_car_schema_without_enum() - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_target_schema(self) -> MetaData: + return get_schema_with_enum_in_sqlalchemy_array_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='{CAR_COLORS_ENUM_NAME}').create(op.get_bind()) - op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', sa.ARRAY(postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}', create_type=False)), nullable=True)) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').create(op.get_bind()) + op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', sa.ARRAY(postgresql.ENUM({', '.join(map(repr, self.new_enum_variants))}, name='{CAR_COLORS_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('{CAR_TABLE_NAME}', '{CAR_COLORS_COLUMN_NAME}') - sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').drop(op.get_bind()) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').drop(op.get_bind()) # ### end Alembic commands ### - """, - ) + """ diff --git a/tests/test_enum_creation/test_create_schema.py b/tests/test_enum_creation/test_create_schema.py index 0261e5f..df499e4 100644 --- a/tests/test_enum_creation/test_create_schema.py +++ b/tests/test_enum_creation/test_create_schema.py @@ -1,39 +1,40 @@ -import textwrap -from typing import TYPE_CHECKING - -from alembic import autogenerate +from sqlalchemy import Table, Column, Integer, MetaData from sqlalchemy.dialects import postgresql +from tests.base.run_migration_test_abc import CompareAndRunTestCase from tests.schemas import ( USER_TABLE_NAME, USER_STATUS_ENUM_NAME, USER_STATUS_COLUMN_NAME, ) -from tests.utils.migration_context import create_migration_context - -if TYPE_CHECKING: - from sqlalchemy import Connection -from sqlalchemy import Table, Column, Integer, MetaData -def test_create_enum_on_create_table_inside_new_schema(connection: "Connection"): +class TestCreateEnumOnCreateTableInsideNewSchema(CompareAndRunTestCase): """Check that library correctly creates enum before its use inside create_table inside new schema""" + + disable_running = True + new_enum_variants = ["active", "passive"] non_existing_schema = "non_existing_schema" - target_schema = MetaData(schema=non_existing_schema) - - Table( - USER_TABLE_NAME, - target_schema, - Column("id", Integer, primary_key=True), - Column( - USER_STATUS_COLUMN_NAME, - postgresql.ENUM(*new_enum_variants, name=USER_STATUS_ENUM_NAME, metadata=target_schema), - ), - ) + def get_database_schema(self) -> MetaData: + return MetaData() + + def get_target_schema(self) -> MetaData: + target_schema = MetaData(schema=self.non_existing_schema) + Table( + USER_TABLE_NAME, + target_schema, + Column("id", Integer, primary_key=True), + Column( + USER_STATUS_COLUMN_NAME, + postgresql.ENUM(*self.new_enum_variants, name=USER_STATUS_ENUM_NAME, metadata=target_schema), + ), + ) + return target_schema - expected_upgrade = f""" + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### sa.Enum('active', 'passive', name='user_status', schema='non_existing_schema').create(op.get_bind()) op.create_table('users', @@ -43,25 +44,12 @@ def test_create_enum_on_create_table_inside_new_schema(connection: "Connection") schema='non_existing_schema' ) # ### end Alembic commands ### - """ - expected_downgrade = f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('users', schema='non_existing_schema') sa.Enum('active', 'passive', name='user_status', schema='non_existing_schema').drop(op.get_bind()) # ### end Alembic commands ### - """ - - migration_context = create_migration_context(connection, target_schema) - - template_args = {} - # todo _render_migration_diffs marked as legacy, maybe find something else - autogenerate._render_migration_diffs(migration_context, template_args) - - upgrade_code = textwrap.dedent(" " + template_args["upgrades"]) - downgrade_code = textwrap.dedent(" " + template_args["downgrades"]) - - expected_upgrade = textwrap.dedent(expected_upgrade).strip("\n ") - expected_downgrade = textwrap.dedent(expected_downgrade).strip("\n ") - - 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}" + """ diff --git a/tests/test_enum_creation/test_create_table.py b/tests/test_enum_creation/test_create_table.py index b798719..e727f94 100644 --- a/tests/test_enum_creation/test_create_table.py +++ b/tests/test_enum_creation/test_create_table.py @@ -1,9 +1,8 @@ -from typing import TYPE_CHECKING - import sqlalchemy +from sqlalchemy import Table, Column, Integer, MetaData from sqlalchemy.dialects import postgresql -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, @@ -17,21 +16,20 @@ ANOTHER_SCHEMA_NAME, ) -if TYPE_CHECKING: - from sqlalchemy import Connection -from sqlalchemy import Table, Column, Integer, MetaData - -def test_create_enum_on_create_table(connection: "Connection"): +class TestCreateEnumOnCreateTable(CompareAndRunTestCase): """Check that library correctly creates enum before its use inside create_table""" + new_enum_variants = ["active", "passive"] - target_schema = get_schema_with_enum_variants(new_enum_variants) + def get_database_schema(self) -> MetaData: + return MetaData() + + 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('active', 'passive', name='user_status').create(op.get_bind()) op.create_table('{USER_TABLE_NAME}', @@ -40,65 +38,70 @@ def test_create_enum_on_create_table(connection: "Connection"): sa.PrimaryKeyConstraint('id') ) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('{USER_TABLE_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_on_create_table_with_array(connection: "Connection"): +class TestCreateEnumOnCreateTableWithArray(CompareAndRunTestCase): """Check that library correctly creates enum before its use inside create_table. Enum is used in ARRAY""" - database_schema = get_car_schema_without_enum() - database_schema.create_all(connection) new_enum_variants = ["black", "white", "red", "green", "blue", "other"] - target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) + def get_database_schema(self) -> MetaData: + return get_car_schema_without_enum() - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_target_schema(self) -> MetaData: + return get_schema_with_enum_in_array_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='{CAR_COLORS_ENUM_NAME}').create(op.get_bind()) - op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}', create_type=False)), nullable=True)) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').create(op.get_bind()) + op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, self.new_enum_variants))}, name='{CAR_COLORS_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('{CAR_TABLE_NAME}', '{CAR_COLORS_COLUMN_NAME}') - sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').drop(op.get_bind()) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').drop(op.get_bind()) # ### end Alembic commands ### - """, - ) + """ -def test_create_enum_on_create_table_with_sa_enum(connection: "Connection"): +class TestCreateEnumOnCreateTableWithSaEnum(CompareAndRunTestCase): """Check that library correctly creates enum before its use inside create_table. sqlalchemy.Enum is used sqlalchemy.Enum should be converted to postgresql.ENUM to specify create_type=False """ + new_enum_variants = ["active", "passive"] - target_schema = MetaData() - - 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), - ), - ) - - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_database_schema(self) -> MetaData: + return MetaData() + + 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), + ), + ) + return target_schema + + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### sa.Enum('active', 'passive', name='user_status').create(op.get_bind()) op.create_table('{USER_TABLE_NAME}', @@ -107,96 +110,107 @@ def test_create_enum_on_create_table_with_sa_enum(connection: "Connection"): sa.PrimaryKeyConstraint('id') ) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('{USER_TABLE_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_on_create_table_with_another_schema(connection: "Connection"): +class TestCreateEnumOnCreateTableWithAnotherSchema(CompareAndRunTestCase): """Check that library correctly creates enum before its use inside create_table inside another table_schema""" + new_enum_variants = ["active", "passive"] - target_schema = MetaData() - - Table( - USER_TABLE_NAME, - target_schema, - Column("id", Integer, primary_key=True), - Column( - USER_STATUS_COLUMN_NAME, - postgresql.ENUM( - *new_enum_variants, - name=USER_STATUS_ENUM_NAME, - schema=ANOTHER_SCHEMA_NAME, + def get_database_schema(self) -> MetaData: + return MetaData() + + 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, + postgresql.ENUM( + *self.new_enum_variants, + name=USER_STATUS_ENUM_NAME, + schema=ANOTHER_SCHEMA_NAME, + ), ), - ), - schema=ANOTHER_SCHEMA_NAME, - ) - - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + schema=ANOTHER_SCHEMA_NAME, + ) + + return target_schema + + 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}', schema='{ANOTHER_SCHEMA_NAME}').create(op.get_bind()) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}').create(op.get_bind()) op.create_table('{USER_TABLE_NAME}', sa.Column('id', sa.Integer(), nullable=False), - sa.Column('status', postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}', create_type=False), nullable=True), + sa.Column('status', postgresql.ENUM({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}', create_type=False), nullable=True), sa.PrimaryKeyConstraint('id'), schema='{ANOTHER_SCHEMA_NAME}' ) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('{USER_TABLE_NAME}', schema='{ANOTHER_SCHEMA_NAME}') - sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}').drop(op.get_bind()) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}').drop(op.get_bind()) # ### end Alembic commands ### - """, - ) + """ -def test_create_enum_on_create_table_with_another_schema_on_metadata( - connection: "Connection", -): +class TestCreateEnumOnCreateTableWithAnotherSchemaOnMetadata(CompareAndRunTestCase): """Check that library correctly creates enum before its use inside create_table inside another schema, specified on Metadata""" + new_enum_variants = ["active", "passive"] - target_schema = MetaData(schema=ANOTHER_SCHEMA_NAME) - - Table( - USER_TABLE_NAME, - target_schema, - Column("id", Integer, primary_key=True), - Column( - USER_STATUS_COLUMN_NAME, - postgresql.ENUM(*new_enum_variants, name=USER_STATUS_ENUM_NAME, metadata=target_schema), - ), - ) - - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_database_schema(self) -> MetaData: + return MetaData() + + def get_target_schema(self) -> MetaData: + target_schema = MetaData(schema=ANOTHER_SCHEMA_NAME) + + Table( + USER_TABLE_NAME, + target_schema, + Column("id", Integer, primary_key=True), + Column( + USER_STATUS_COLUMN_NAME, + postgresql.ENUM(*self.new_enum_variants, name=USER_STATUS_ENUM_NAME, metadata=target_schema), + ), + ) + + return target_schema + + 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}', schema='{ANOTHER_SCHEMA_NAME}').create(op.get_bind()) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}').create(op.get_bind()) op.create_table('{USER_TABLE_NAME}', sa.Column('id', sa.Integer(), nullable=False), - sa.Column('status', postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}', create_type=False), nullable=True), + sa.Column('status', postgresql.ENUM({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}', create_type=False), nullable=True), sa.PrimaryKeyConstraint('id'), schema='{ANOTHER_SCHEMA_NAME}' ) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('{USER_TABLE_NAME}', schema='{ANOTHER_SCHEMA_NAME}') - sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}').drop(op.get_bind()) + sa.Enum({', '.join(map(repr, self.new_enum_variants))}, name='{USER_STATUS_ENUM_NAME}', schema='{ANOTHER_SCHEMA_NAME}').drop(op.get_bind()) # ### end Alembic commands ### - """, - ) + """ diff --git a/tests/test_enum_creation/test_drop_column.py b/tests/test_enum_creation/test_drop_column.py index 1c8c615..e188247 100644 --- a/tests/test_enum_creation/test_drop_column.py +++ b/tests/test_enum_creation/test_drop_column.py @@ -1,9 +1,10 @@ from typing import TYPE_CHECKING from alembic_postgresql_enum.operations import DropEnumOp +from tests.base.run_migration_test_abc import CompareAndRunTestCase if TYPE_CHECKING: - from sqlalchemy import Connection + from sqlalchemy import Connection, MetaData from tests.schemas import ( get_schema_with_enum_in_array_variants, @@ -18,7 +19,6 @@ from alembic.autogenerate import api from alembic.operations import ops -from tests.base.render_and_run import compare_and_run from tests.schemas import ( get_schema_with_enum_variants, USER_TABLE_NAME, @@ -31,58 +31,64 @@ if TYPE_CHECKING: from sqlalchemy import Connection +from sqlalchemy import MetaData -def test_delete_enum_after_drop_column(connection: "Connection"): +class TestDeleteEnumAfterDropColumn(CompareAndRunTestCase): """Check that library correctly removes unused enum after drop_column""" + enum_variants_to_delete = ["active", "passive"] - database_schema = get_schema_with_enum_variants(enum_variants_to_delete) - database_schema.create_all(connection) - target_schema = get_schema_without_enum() + def get_database_schema(self) -> MetaData: + return get_schema_with_enum_variants(self.enum_variants_to_delete) + + def get_target_schema(self) -> MetaData: + return get_schema_without_enum() - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_expected_upgrade(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, enum_variants_to_delete))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind()) + sa.Enum({', '.join(map(repr, self.enum_variants_to_delete))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind()) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + # For some reason alembic decided to add redundant autoincrement=False on downgrade + return f""" # ### commands auto generated by Alembic - please adjust! ### - sa.Enum({', '.join(map(repr, enum_variants_to_delete))}, 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, enum_variants_to_delete))}, name='{USER_STATUS_ENUM_NAME}', create_type=False), autoincrement=False, nullable=True)) + sa.Enum({', '.join(map(repr, self.enum_variants_to_delete))}, 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.enum_variants_to_delete))}, name='{USER_STATUS_ENUM_NAME}', create_type=False), autoincrement=False, nullable=True)) # ### end Alembic commands ### - """, - ) # For some reason alembic decided to add redundant autoincrement=False on downgrade + """ -def test_delete_enum_after_drop_column_with_array(connection: "Connection"): +class TestDeleteEnumAfterDropColumnWithArray(CompareAndRunTestCase): """Check that library correctly removes unused enum after drop_column. Enum is used in ARRAY""" + enum_variants_to_delete = ["black", "white", "red", "green", "blue", "other"] - database_schema = get_schema_with_enum_in_array_variants(enum_variants_to_delete) - database_schema.create_all(connection) - target_schema = get_car_schema_without_enum() + def get_database_schema(self) -> MetaData: + return get_schema_with_enum_in_array_variants(self.enum_variants_to_delete) - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_target_schema(self) -> MetaData: + return get_car_schema_without_enum() + + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_column('{CAR_TABLE_NAME}', '{CAR_COLORS_COLUMN_NAME}') - sa.Enum({', '.join(map(repr, enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}').drop(op.get_bind()) + sa.Enum({', '.join(map(repr, self.enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}').drop(op.get_bind()) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### - sa.Enum({', '.join(map(repr, enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}').create(op.get_bind()) - op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}', create_type=False)), autoincrement=False, nullable=True)) + sa.Enum({', '.join(map(repr, self.enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}').create(op.get_bind()) + op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, self.enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}', create_type=False)), autoincrement=False, nullable=True)) # ### end Alembic commands ### - """, - ) + """ def test_delete_enum_diff_tuple(connection: "Connection"): diff --git a/tests/test_enum_creation/test_drop_table.py b/tests/test_enum_creation/test_drop_table.py index 9af2b39..22ccf74 100644 --- a/tests/test_enum_creation/test_drop_table.py +++ b/tests/test_enum_creation/test_drop_table.py @@ -1,36 +1,34 @@ -from typing import TYPE_CHECKING - from sqlalchemy import MetaData -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, USER_STATUS_ENUM_NAME, ) -if TYPE_CHECKING: - from sqlalchemy import Connection - -def test_drop_enum_after_drop_table(connection: "Connection"): +class TestDropEnumAfterDropTable(CompareAndRunTestCase): """Check that library correctly drop enum after drop_table""" + dropped_enum_variants = ["active", "passive"] - database_schema = get_schema_with_enum_variants(dropped_enum_variants) - database_schema.create_all(connection) - target_schema = MetaData() + def get_database_schema(self) -> MetaData: + return get_schema_with_enum_variants(self.dropped_enum_variants) - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_target_schema(self) -> MetaData: + return MetaData() + + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('{USER_TABLE_NAME}') - sa.Enum({', '.join(map(repr, dropped_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind()) + sa.Enum({', '.join(map(repr, self.dropped_enum_variants))}, name='{USER_STATUS_ENUM_NAME}').drop(op.get_bind()) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### sa.Enum('active', 'passive', name='user_status').create(op.get_bind()) op.create_table('users', @@ -39,5 +37,4 @@ def test_drop_enum_after_drop_table(connection: "Connection"): sa.PrimaryKeyConstraint('id', name='users_pkey') ) # ### end Alembic commands ### - """, - ) + """ diff --git a/tests/test_without_schema_changes/test_explicit_schema.py b/tests/test_without_schema_changes/test_explicit_schema.py index 0afad2b..614eec6 100644 --- a/tests/test_without_schema_changes/test_explicit_schema.py +++ b/tests/test_without_schema_changes/test_explicit_schema.py @@ -7,7 +7,7 @@ get_defined_enums, get_declared_enums, ) -from tests.base.render_and_run import compare_and_run +from tests.base.run_migration_test_abc import CompareAndRunTestCase from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA if TYPE_CHECKING: @@ -68,23 +68,23 @@ def test_get_declared_enums(connection: "Connection"): } -def test_compare_and_run(connection: "Connection"): - database_schema = my_metadata - database_schema.create_all(connection) +class TestCompareAndRun(CompareAndRunTestCase): + def get_database_schema(self) -> MetaData: + return my_metadata - target_schema = my_metadata + def get_target_schema(self) -> MetaData: + return my_metadata - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - ) + """ diff --git a/tests/test_without_schema_changes/test_implicit_schema_1_4_style.py b/tests/test_without_schema_changes/test_implicit_schema_1_4_style.py index 27d870a..83214b7 100644 --- a/tests/test_without_schema_changes/test_implicit_schema_1_4_style.py +++ b/tests/test_without_schema_changes/test_implicit_schema_1_4_style.py @@ -7,7 +7,7 @@ get_defined_enums, get_declared_enums, ) -from tests.base.render_and_run import compare_and_run +from tests.base.run_migration_test_abc import CompareAndRunTestCase from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA if TYPE_CHECKING: @@ -80,13 +80,15 @@ def test_get_declared_enums(connection: "Connection"): } -def test_compare_and_run_create_table(connection: "Connection"): - target_schema = my_metadata +class TestCompareAndRunCreateTable(CompareAndRunTestCase): + def get_database_schema(self) -> MetaData: + return MetaData() - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_target_schema(self) -> MetaData: + return my_metadata + + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### sa.Enum('PENDING', 'SUCCESS', 'FAILED', name='test_status').create(op.get_bind()) op.create_table('test', @@ -96,33 +98,34 @@ def test_compare_and_run_create_table(connection: "Connection"): schema='another' ) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('test', schema='another') sa.Enum('PENDING', 'SUCCESS', 'FAILED', name='test_status').drop(op.get_bind()) # ### end Alembic commands ### - """, - ) + """ -def test_compare_and_run(connection: "Connection"): - database_schema = my_metadata - database_schema.create_all(connection) +class TestCompareAndRun(CompareAndRunTestCase): + def get_database_schema(self) -> MetaData: + return my_metadata - target_schema = my_metadata + def get_target_schema(self) -> MetaData: + return my_metadata - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - ) + """ diff --git a/tests/test_without_schema_changes/test_implicit_schema_2_0_style.py b/tests/test_without_schema_changes/test_implicit_schema_2_0_style.py index 50750ae..9f1b4f9 100644 --- a/tests/test_without_schema_changes/test_implicit_schema_2_0_style.py +++ b/tests/test_without_schema_changes/test_implicit_schema_2_0_style.py @@ -10,7 +10,7 @@ get_defined_enums, get_declared_enums, ) -from tests.base.render_and_run import compare_and_run +from tests.base.run_migration_test_abc import CompareAndRunTestCase from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA if TYPE_CHECKING: @@ -81,14 +81,15 @@ def test_get_declared_enums(connection: "Connection"): @pytest.mark.skipif(sqlalchemy.__version__.startswith("1."), reason="Table are made in 2.0 style") -def test_compare_and_run_create_table(connection: "Connection"): - my_metadata = get_my_metadata() - target_schema = my_metadata +class TestCompareAndRunCreateTable(CompareAndRunTestCase): + def get_database_schema(self) -> MetaData: + return MetaData() + + def get_target_schema(self) -> MetaData: + return get_my_metadata() - 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('PENDING', 'SUCCESS', 'FAILED', name='test_status').create(op.get_bind()) op.create_table('test', @@ -98,35 +99,35 @@ def test_compare_and_run_create_table(connection: "Connection"): schema='another' ) # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### op.drop_table('test', schema='another') sa.Enum('PENDING', 'SUCCESS', 'FAILED', name='test_status').drop(op.get_bind()) # ### end Alembic commands ### - """, - ) + """ @pytest.mark.skipif(sqlalchemy.__version__.startswith("1."), reason="Table are made in 2.0 style") -def test_compare_and_run(connection: "Connection"): - my_metadata = get_my_metadata() - database_schema = my_metadata - database_schema.create_all(connection) +class TestCompareAndRun(CompareAndRunTestCase): + def get_database_schema(self) -> MetaData: + return get_my_metadata() - target_schema = my_metadata + def get_target_schema(self) -> MetaData: + return get_my_metadata() - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - ) + """ diff --git a/tests/test_without_schema_changes/test_inherit_schema_true.py b/tests/test_without_schema_changes/test_inherit_schema_true.py index 298aeca..9b0578f 100644 --- a/tests/test_without_schema_changes/test_inherit_schema_true.py +++ b/tests/test_without_schema_changes/test_inherit_schema_true.py @@ -7,7 +7,7 @@ get_defined_enums, get_declared_enums, ) -from tests.base.render_and_run import compare_and_run +from tests.base.run_migration_test_abc import CompareAndRunTestCase from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA if TYPE_CHECKING: @@ -68,23 +68,23 @@ def test_get_declared_enums(connection: "Connection"): } -def test_compare_and_run(connection: "Connection"): - database_schema = my_metadata - database_schema.create_all(connection) +class TestCompareAndRun(CompareAndRunTestCase): + def get_database_schema(self) -> MetaData: + return my_metadata - target_schema = my_metadata + def get_target_schema(self) -> MetaData: + return my_metadata - compare_and_run( - connection, - target_schema, - expected_upgrade=f""" + def get_expected_upgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - expected_downgrade=f""" + """ + + def get_expected_downgrade(self) -> str: + return f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, - ) + """