From 05221ff938e0e32c7528d7a518e22e22250a0404 Mon Sep 17 00:00:00 2001 From: rusty Date: Sat, 6 Jan 2024 11:23:34 +0400 Subject: [PATCH] Reformat tests files with black --- tests/base/render_and_run.py | 62 +++--- tests/fixtures/db.py | 11 +- .../get_enum_data/test_get_declared_enums.py | 31 ++- tests/get_enum_data/test_get_defined_enums.py | 17 +- .../test_type_decorator_enums.py | 53 +++-- tests/schemas.py | 109 ++++++---- tests/sync_enum_values/test_array_column.py | 77 ++++--- .../test_rename_default_if_required.py | 34 ++- tests/sync_enum_values/test_render.py | 65 +++--- .../test_run_in_different_schema.py | 16 +- .../test_run_sync_enum_values.py | 200 +++++++++++------- tests/test_alter_column/test_text_column.py | 43 ++-- tests/test_constraints/test_unique.py | 76 ++++--- tests/test_enum_creation/test_add_column.py | 75 ++++--- .../test_enum_creation/test_create_schema.py | 37 ++-- tests/test_enum_creation/test_create_table.py | 112 +++++++--- tests/test_enum_creation/test_drop_column.py | 64 ++++-- tests/test_enum_creation/test_drop_table.py | 19 +- .../test_explicit_schema.py | 35 ++- .../test_implicit_schema_1_4_style.py | 55 +++-- .../test_implicit_schema_2_0_style.py | 68 ++++-- .../test_inherit_schema_true.py | 35 ++- tests/utils/migration_context.py | 4 +- 23 files changed, 837 insertions(+), 461 deletions(-) diff --git a/tests/base/render_and_run.py b/tests/base/render_and_run.py index 96fefa3..332b428 100644 --- a/tests/base/render_and_run.py +++ b/tests/base/render_and_run.py @@ -14,9 +14,13 @@ from sqlalchemy import Connection -def compare_and_run(connection: 'Connection', target_schema: Union[MetaData, List[MetaData]], *, - expected_upgrade: str, - expected_downgrade: str): +def compare_and_run( + connection: "Connection", + target_schema: Union[MetaData, List[MetaData]], + *, + expected_upgrade: str, + expected_downgrade: str, +): """Compares generated migration script is equal to expected_upgrade and expected_downgrade, then runs it""" migration_context = create_migration_context(connection, target_schema) @@ -26,24 +30,34 @@ def compare_and_run(connection: 'Connection', target_schema: Union[MetaData, Lis # 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}' - - exec(upgrade_code, { # todo Use imports from template_args - 'op': op, - 'sa': sqlalchemy, - 'postgresql': postgresql, - 'ColumnType': ColumnType, - }) - exec(downgrade_code, { - 'op': op, - 'sa': sqlalchemy, - 'postgresql': postgresql, - 'ColumnType': ColumnType, - }) + 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}" + + exec( + upgrade_code, + { # todo Use imports from template_args + "op": op, + "sa": sqlalchemy, + "postgresql": postgresql, + "ColumnType": ColumnType, + }, + ) + exec( + downgrade_code, + { + "op": op, + "sa": sqlalchemy, + "postgresql": postgresql, + "ColumnType": ColumnType, + }, + ) diff --git a/tests/fixtures/db.py b/tests/fixtures/db.py index c46725f..7b79703 100644 --- a/tests/fixtures/db.py +++ b/tests/fixtures/db.py @@ -9,20 +9,25 @@ try: import dotenv + dotenv.load_dotenv() except ImportError: pass -database_uri = os.getenv('DATABASE_URI') +database_uri = os.getenv("DATABASE_URI") @pytest.fixture def connection() -> Generator: engine = create_engine(database_uri) with engine.connect() as conn: - conn.execute(sqlalchemy.text(f''' + conn.execute( + sqlalchemy.text( + f""" DROP SCHEMA {DEFAULT_SCHEMA} CASCADE; CREATE SCHEMA {DEFAULT_SCHEMA}; DROP SCHEMA IF EXISTS {ANOTHER_SCHEMA_NAME} CASCADE; CREATE SCHEMA {ANOTHER_SCHEMA_NAME}; - ''')) + """ + ) + ) yield conn diff --git a/tests/get_enum_data/test_get_declared_enums.py b/tests/get_enum_data/test_get_declared_enums.py index 68f3a3a..0458837 100644 --- a/tests/get_enum_data/test_get_declared_enums.py +++ b/tests/get_enum_data/test_get_declared_enums.py @@ -1,31 +1,42 @@ from typing import TYPE_CHECKING from alembic_postgresql_enum.get_enum_data import TableReference, get_declared_enums -from tests.schemas import get_schema_with_enum_variants, DEFAULT_SCHEMA, USER_STATUS_ENUM_NAME, USER_TABLE_NAME, \ - USER_STATUS_COLUMN_NAME, get_schema_by_declared_enum_values, get_declared_enum_values_with_orders_and_users +from tests.schemas import ( + get_schema_with_enum_variants, + DEFAULT_SCHEMA, + USER_STATUS_ENUM_NAME, + USER_TABLE_NAME, + USER_STATUS_COLUMN_NAME, + get_schema_by_declared_enum_values, + get_declared_enum_values_with_orders_and_users, +) if TYPE_CHECKING: from sqlalchemy import Connection -def test_with_user_schema(connection: 'Connection'): +def test_with_user_schema(connection: "Connection"): enum_variants = ["active", "passive"] declared_schema = get_schema_with_enum_variants(enum_variants) - function_result = get_declared_enums(declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection + ) - assert function_result.enum_values == { - USER_STATUS_ENUM_NAME: tuple(enum_variants) - } + assert function_result.enum_values == {USER_STATUS_ENUM_NAME: tuple(enum_variants)} assert function_result.enum_table_references == { - USER_STATUS_ENUM_NAME: frozenset((TableReference(USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),)) + USER_STATUS_ENUM_NAME: frozenset( + (TableReference(USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),) + ) } -def test_with_multiple_enums(connection: 'Connection'): +def test_with_multiple_enums(connection: "Connection"): declared_enum_values = get_declared_enum_values_with_orders_and_users() declared_schema = get_schema_by_declared_enum_values(declared_enum_values) - function_result = get_declared_enums(declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection + ) assert function_result == declared_enum_values diff --git a/tests/get_enum_data/test_get_defined_enums.py b/tests/get_enum_data/test_get_defined_enums.py index 32dceec..88b8338 100644 --- a/tests/get_enum_data/test_get_defined_enums.py +++ b/tests/get_enum_data/test_get_defined_enums.py @@ -4,23 +4,26 @@ from sqlalchemy import Connection from alembic_postgresql_enum.get_enum_data import get_defined_enums -from tests.schemas import get_schema_with_enum_variants, DEFAULT_SCHEMA, USER_STATUS_ENUM_NAME, \ - get_schema_by_declared_enum_values, get_declared_enum_values_with_orders_and_users +from tests.schemas import ( + get_schema_with_enum_variants, + DEFAULT_SCHEMA, + USER_STATUS_ENUM_NAME, + get_schema_by_declared_enum_values, + get_declared_enum_values_with_orders_and_users, +) -def test_get_defined_enums(connection: 'Connection'): +def test_get_defined_enums(connection: "Connection"): enum_variants = ["active", "passive"] defined_schema = get_schema_with_enum_variants(enum_variants) defined_schema.create_all(connection) function_result = get_defined_enums(connection, DEFAULT_SCHEMA) - assert function_result == { - USER_STATUS_ENUM_NAME: tuple(enum_variants) - } + assert function_result == {USER_STATUS_ENUM_NAME: tuple(enum_variants)} -def test_with_multiple_enums(connection: 'Connection'): +def test_with_multiple_enums(connection: "Connection"): declared_enum_values = get_declared_enum_values_with_orders_and_users() defined_schema = get_schema_by_declared_enum_values(declared_enum_values) diff --git a/tests/get_enum_data/test_type_decorator_enums.py b/tests/get_enum_data/test_type_decorator_enums.py index 19b5a7d..37759dc 100644 --- a/tests/get_enum_data/test_type_decorator_enums.py +++ b/tests/get_enum_data/test_type_decorator_enums.py @@ -15,6 +15,7 @@ class ValuesEnum(sqlalchemy.types.TypeDecorator): """Custom enum wrapper that forces columns to store enum values, not names""" + impl = sqlalchemy.types.Enum cache_ok = True @@ -42,44 +43,54 @@ def process(value): return process -ORDER_TABLE_NAME = 'order' -ORDER_DELIVERY_STATUS_COLUMN_NAME = 'delivery_status' -ORDER_DELIVERY_STATUS_ENUM_NAME = 'order_delivery_status' +ORDER_TABLE_NAME = "order" +ORDER_DELIVERY_STATUS_COLUMN_NAME = "delivery_status" +ORDER_DELIVERY_STATUS_ENUM_NAME = "order_delivery_status" class OrderDeliveryStatus(enum.Enum): - WAITING_FOR_WORKER = 'waiting_for_worker' - WAITING_FOR_WORKER_TO_ARRIVE = 'waiting_for_worker_to_arrive' - WORKER_ARRIVED = 'worker_arrived' - IN_PROGRESS = 'in_progress' - WAITING_FOR_APPROVAL = 'waiting_for_approval' - DISAPPROVED = 'disapproved' - DONE = 'done' - REFUNDED = 'refunded' - BANNED = 'banned' - CANCELED = 'canceled' + WAITING_FOR_WORKER = "waiting_for_worker" + WAITING_FOR_WORKER_TO_ARRIVE = "waiting_for_worker_to_arrive" + WORKER_ARRIVED = "worker_arrived" + IN_PROGRESS = "in_progress" + WAITING_FOR_APPROVAL = "waiting_for_approval" + DISAPPROVED = "disapproved" + DONE = "done" + REFUNDED = "refunded" + BANNED = "banned" + CANCELED = "canceled" def get_schema_with_custom_enum() -> MetaData: schema = MetaData() - Table(ORDER_TABLE_NAME, - schema, - Column(ORDER_DELIVERY_STATUS_COLUMN_NAME, ValuesEnum(OrderDeliveryStatus, name=ORDER_DELIVERY_STATUS_ENUM_NAME)) - ) + Table( + ORDER_TABLE_NAME, + schema, + Column( + ORDER_DELIVERY_STATUS_COLUMN_NAME, + ValuesEnum(OrderDeliveryStatus, name=ORDER_DELIVERY_STATUS_ENUM_NAME), + ), + ) return schema -def test_get_declared_enums_for_custom_enum(connection: 'Connection'): +def test_get_declared_enums_for_custom_enum(connection: "Connection"): declared_schema = get_schema_with_custom_enum() - function_result = get_declared_enums(declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection + ) assert function_result.enum_values == { # All declared enum variants must be taken from OrderDeliveryStatus values, see ValuesEnum - ORDER_DELIVERY_STATUS_ENUM_NAME: tuple(enum_item.value for enum_item in OrderDeliveryStatus) + ORDER_DELIVERY_STATUS_ENUM_NAME: tuple( + enum_item.value for enum_item in OrderDeliveryStatus + ) } assert function_result.enum_table_references == { - ORDER_DELIVERY_STATUS_ENUM_NAME: frozenset((TableReference(ORDER_TABLE_NAME, ORDER_DELIVERY_STATUS_COLUMN_NAME),)) + ORDER_DELIVERY_STATUS_ENUM_NAME: frozenset( + (TableReference(ORDER_TABLE_NAME, ORDER_DELIVERY_STATUS_COLUMN_NAME),) + ) } diff --git a/tests/schemas.py b/tests/schemas.py index 1dbea51..3314ecd 100644 --- a/tests/schemas.py +++ b/tests/schemas.py @@ -4,18 +4,22 @@ from sqlalchemy import MetaData, Table, Column, Integer from sqlalchemy.dialects import postgresql -from alembic_postgresql_enum.get_enum_data import DeclaredEnumValues, TableReference, ColumnType +from alembic_postgresql_enum.get_enum_data import ( + DeclaredEnumValues, + TableReference, + ColumnType, +) -DEFAULT_SCHEMA = 'public' -USER_TABLE_NAME = 'users' -USER_STATUS_COLUMN_NAME = 'status' -USER_STATUS_ENUM_NAME = 'user_status' +DEFAULT_SCHEMA = "public" +USER_TABLE_NAME = "users" +USER_STATUS_COLUMN_NAME = "status" +USER_STATUS_ENUM_NAME = "user_status" -CAR_TABLE_NAME = 'cars' -CAR_COLORS_COLUMN_NAME = 'colors' -CAR_COLORS_ENUM_NAME = 'car_color' +CAR_TABLE_NAME = "cars" +CAR_COLORS_COLUMN_NAME = "colors" +CAR_COLORS_ENUM_NAME = "car_color" -ANOTHER_SCHEMA_NAME = 'another' +ANOTHER_SCHEMA_NAME = "another" def get_schema_with_enum_variants(variants: List[str]) -> MetaData: @@ -25,7 +29,10 @@ def get_schema_with_enum_variants(variants: List[str]) -> MetaData: USER_TABLE_NAME, schema, Column("id", Integer, primary_key=True), - Column(USER_STATUS_COLUMN_NAME, postgresql.ENUM(*variants, name=USER_STATUS_ENUM_NAME)) + Column( + USER_STATUS_COLUMN_NAME, + postgresql.ENUM(*variants, name=USER_STATUS_ENUM_NAME), + ), ) return schema @@ -38,7 +45,10 @@ def get_schema_with_enum_in_array_variants(variants: List[str]) -> MetaData: CAR_TABLE_NAME, schema, Column("id", Integer, primary_key=True), - Column(CAR_COLORS_COLUMN_NAME, postgresql.ARRAY(postgresql.ENUM(*variants, name=CAR_COLORS_ENUM_NAME))) + Column( + CAR_COLORS_COLUMN_NAME, + postgresql.ARRAY(postgresql.ENUM(*variants, name=CAR_COLORS_ENUM_NAME)), + ), ) return schema @@ -55,6 +65,7 @@ def get_schema_without_enum() -> MetaData: return schema + def get_car_schema_without_enum() -> MetaData: schema = MetaData() @@ -70,50 +81,66 @@ def get_car_schema_without_enum() -> MetaData: def get_declared_enum_values_with_orders_and_users() -> DeclaredEnumValues: return DeclaredEnumValues( enum_values={ - 'user_status_enum': ('active', 'inactive', 'banned'), - 'order_status_enum': ( - 'waiting_for_worker', - 'waiting_for_worker_to_arrive', - 'worker_arrived', - 'in_progress', - 'waiting_for_approval', - 'disapproved', - 'done', - 'refunded', - 'banned', - 'canceled', + "user_status_enum": ("active", "inactive", "banned"), + "order_status_enum": ( + "waiting_for_worker", + "waiting_for_worker_to_arrive", + "worker_arrived", + "in_progress", + "waiting_for_approval", + "disapproved", + "done", + "refunded", + "banned", + "canceled", ), - 'car_color_enum': ('black', 'white', 'red', 'green', 'blue', 'other') + "car_color_enum": ("black", "white", "red", "green", "blue", "other"), }, enum_table_references={ - 'user_status_enum': frozenset(( - TableReference('users', 'user_status'), - TableReference('users', 'last_user_status'), - TableReference('orders', 'user_status'), - )), - 'order_status_enum': frozenset(( - TableReference('orders', 'order_status'), - )), - 'car_color_enum': frozenset(( - TableReference('cars', 'colors', ColumnType.ARRAY), - )) - } + "user_status_enum": frozenset( + ( + TableReference("users", "user_status"), + TableReference("users", "last_user_status"), + TableReference("orders", "user_status"), + ) + ), + "order_status_enum": frozenset((TableReference("orders", "order_status"),)), + "car_color_enum": frozenset( + (TableReference("cars", "colors", ColumnType.ARRAY),) + ), + }, ) -def _enum_column_factory(target: DeclaredEnumValues, column_name: str, enum_name: str, column_type: ColumnType) -> Column: +def _enum_column_factory( + target: DeclaredEnumValues, + column_name: str, + enum_name: str, + column_type: ColumnType, +) -> Column: if column_type == ColumnType.COMMON: - return Column(column_name, postgresql.ENUM(*target.enum_values[enum_name], name=enum_name)) - return Column(column_name, column_type.value(postgresql.ENUM(*target.enum_values[enum_name], name=enum_name))) + return Column( + column_name, postgresql.ENUM(*target.enum_values[enum_name], name=enum_name) + ) + return Column( + column_name, + column_type.value( + postgresql.ENUM(*target.enum_values[enum_name], name=enum_name) + ), + ) def get_schema_by_declared_enum_values(target: DeclaredEnumValues) -> MetaData: schema = MetaData() - tables_to_columns: DefaultDict[Any, Set[Tuple[str, str, ColumnType]]] = defaultdict(set) + tables_to_columns: DefaultDict[Any, Set[Tuple[str, str, ColumnType]]] = defaultdict( + set + ) for enum_name, references in target.enum_table_references.items(): for reference in references: - tables_to_columns[reference.table_name].add((reference.column_name, enum_name, reference.column_type)) + tables_to_columns[reference.table_name].add( + (reference.column_name, enum_name, reference.column_type) + ) for table_name, columns_with_enum_names in tables_to_columns.items(): Table( diff --git a/tests/sync_enum_values/test_array_column.py b/tests/sync_enum_values/test_array_column.py index d9f737e..358b72d 100644 --- a/tests/sync_enum_values/test_array_column.py +++ b/tests/sync_enum_values/test_array_column.py @@ -11,24 +11,25 @@ if TYPE_CHECKING: from sqlalchemy import Connection -from tests.schemas import (get_schema_with_enum_in_array_variants, - DEFAULT_SCHEMA, - CAR_TABLE_NAME, - CAR_COLORS_COLUMN_NAME, - CAR_COLORS_ENUM_NAME - ) +from tests.schemas import ( + get_schema_with_enum_in_array_variants, + DEFAULT_SCHEMA, + CAR_TABLE_NAME, + CAR_COLORS_COLUMN_NAME, + CAR_COLORS_ENUM_NAME, +) from tests.utils.migration_context import create_migration_context -def test_add_new_enum_value_render_with_array(connection: 'Connection'): +def test_add_new_enum_value_render_with_array(connection: "Connection"): """Check that enum variants are updated when new variant is added""" - old_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'other'] + old_enum_variants = ["black", "white", "red", "green", "blue", "other"] database_schema = get_schema_with_enum_in_array_variants(old_enum_variants) database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.append('violet') + new_enum_variants.append("violet") target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) @@ -37,32 +38,38 @@ def test_add_new_enum_value_render_with_array(connection: 'Connection'): template_args = {} autogenerate._render_migration_diffs(context, template_args) - assert template_args['imports'] == ('from alembic_postgresql_enum import ColumnType' - '\nfrom alembic_postgresql_enum import TableReference') + assert template_args["imports"] == ( + "from alembic_postgresql_enum import ColumnType" + "\nfrom alembic_postgresql_enum import TableReference" + ) - assert (template_args["upgrades"] == - f"""# ### commands auto generated by Alembic - please adjust! ### + assert ( + template_args["upgrades"] + == f"""# ### commands auto generated by Alembic - please adjust! ### op.sync_enum_values('{DEFAULT_SCHEMA}', '{CAR_COLORS_ENUM_NAME}', [{', '.join(map(repr, new_enum_variants))}], [TableReference('{CAR_TABLE_NAME}', '{CAR_COLORS_COLUMN_NAME}', column_type=ColumnType.ARRAY)], enum_values_to_rename=[]) - # ### end Alembic commands ###""") - assert (template_args["downgrades"] == - f"""# ### commands auto generated by Alembic - please adjust! ### + # ### end Alembic commands ###""" + ) + assert ( + template_args["downgrades"] + == f"""# ### commands auto generated by Alembic - please adjust! ### op.sync_enum_values('{DEFAULT_SCHEMA}', '{CAR_COLORS_ENUM_NAME}', [{', '.join(map(repr, old_enum_variants))}], [TableReference('{CAR_TABLE_NAME}', '{CAR_COLORS_COLUMN_NAME}', column_type=ColumnType.ARRAY)], enum_values_to_rename=[]) - # ### end Alembic commands ###""") + # ### end Alembic commands ###""" + ) -def test_add_new_enum_value_diff_tuple_with_array(connection: 'Connection'): +def test_add_new_enum_value_diff_tuple_with_array(connection: "Connection"): """Check that enum variants are updated when new variant is added""" - old_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'other'] + old_enum_variants = ["black", "white", "red", "green", "blue", "other"] database_schema = get_schema_with_enum_in_array_variants(old_enum_variants) database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.append('violet') + new_enum_variants.append("violet") target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) @@ -81,19 +88,19 @@ def test_add_new_enum_value_diff_tuple_with_array(connection: 'Connection'): SyncEnumValuesOp.operation_name, old_enum_variants, new_enum_variants, - [TableReference(CAR_TABLE_NAME, CAR_COLORS_COLUMN_NAME, ColumnType.ARRAY)] + [TableReference(CAR_TABLE_NAME, CAR_COLORS_COLUMN_NAME, ColumnType.ARRAY)], ) -def test_remove_enum_value_diff_tuple_with_array(connection: 'Connection'): +def test_remove_enum_value_diff_tuple_with_array(connection: "Connection"): """Check that enum variants are updated when new variant is removed""" - old_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'violet', 'other'] + old_enum_variants = ["black", "white", "red", "green", "blue", "violet", "other"] database_schema = get_schema_with_enum_in_array_variants(old_enum_variants) database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('violet') + new_enum_variants.remove("violet") target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) @@ -108,7 +115,12 @@ def test_remove_enum_value_diff_tuple_with_array(connection: 'Connection'): assert len(diffs) == 1 change_variants_diff_tuple = diffs[0] - operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple + ( + operation_name, + old_values, + new_values, + affected_columns, + ) = change_variants_diff_tuple assert operation_name == SyncEnumValuesOp.operation_name assert old_values == old_enum_variants @@ -118,16 +130,16 @@ def test_remove_enum_value_diff_tuple_with_array(connection: 'Connection'): ] -def test_rename_enum_value_diff_tuple_with_array(connection: 'Connection'): +def test_rename_enum_value_diff_tuple_with_array(connection: "Connection"): """Check that enum variants are updated when a variant is renamed""" - old_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'other'] + old_enum_variants = ["black", "white", "red", "green", "blue", "other"] database_schema = get_schema_with_enum_in_array_variants(old_enum_variants) database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('green') - new_enum_variants.append('violet') + new_enum_variants.remove("green") + new_enum_variants.append("violet") target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) @@ -142,7 +154,12 @@ def test_rename_enum_value_diff_tuple_with_array(connection: 'Connection'): assert len(diffs) == 1 change_variants_diff_tuple = diffs[0] - operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple + ( + operation_name, + old_values, + new_values, + affected_columns, + ) = change_variants_diff_tuple assert operation_name == SyncEnumValuesOp.operation_name assert old_values == old_enum_variants diff --git a/tests/sync_enum_values/test_rename_default_if_required.py b/tests/sync_enum_values/test_rename_default_if_required.py index 9c81ebc..1a002ee 100644 --- a/tests/sync_enum_values/test_rename_default_if_required.py +++ b/tests/sync_enum_values/test_rename_default_if_required.py @@ -1,31 +1,43 @@ -from alembic_postgresql_enum.sql_commands.column_default import rename_default_if_required +from alembic_postgresql_enum.sql_commands.column_default import ( + rename_default_if_required, +) def test_without_renames_with_schema(): old_default_value = "'passive'::test.order_status_old" - assert (rename_default_if_required('test', old_default_value, 'order_status', []) - == "'passive'::test.order_status") + assert ( + rename_default_if_required("test", old_default_value, "order_status", []) + == "'passive'::test.order_status" + ) def test_with_renames_with_schema(): old_default_value = "'passive'::test.order_status_old" - assert (rename_default_if_required('test', old_default_value, 'order_status', [ - ('passive', 'inactive') - ]) == "'inactive'::test.order_status") + assert ( + rename_default_if_required( + "test", old_default_value, "order_status", [("passive", "inactive")] + ) + == "'inactive'::test.order_status" + ) def test_array_with_renames_with_schema(): old_default_value = """'{"passive"}'::test.order_status_old""" - assert (rename_default_if_required('test', old_default_value, 'order_status', [ - ('passive', 'inactive') - ]) == """'{"inactive"}'::test.order_status""") + assert ( + rename_default_if_required( + "test", old_default_value, "order_status", [("passive", "inactive")] + ) + == """'{"inactive"}'::test.order_status""" + ) def test_array_default_value_with_schema(): old_default_value = """'{}'::test.order_status_old[]""" - assert (rename_default_if_required('test', old_default_value, 'order_status', []) - == """'{}'::test.order_status[]""") + assert ( + rename_default_if_required("test", old_default_value, "order_status", []) + == """'{}'::test.order_status[]""" + ) diff --git a/tests/sync_enum_values/test_render.py b/tests/sync_enum_values/test_render.py index f3da8f8..305d952 100644 --- a/tests/sync_enum_values/test_render.py +++ b/tests/sync_enum_values/test_render.py @@ -10,16 +10,17 @@ if TYPE_CHECKING: from sqlalchemy import Connection -from tests.schemas import (get_schema_with_enum_variants, - USER_TABLE_NAME, - USER_STATUS_ENUM_NAME, - USER_STATUS_COLUMN_NAME, - DEFAULT_SCHEMA - ) +from tests.schemas import ( + get_schema_with_enum_variants, + USER_TABLE_NAME, + USER_STATUS_ENUM_NAME, + USER_STATUS_COLUMN_NAME, + DEFAULT_SCHEMA, +) from tests.utils.migration_context import create_migration_context -def test_add_new_enum_value_render(connection: 'Connection'): +def test_add_new_enum_value_render(connection: "Connection"): """Check that enum variants are updated when new variant is added""" old_enum_variants = ["active", "passive"] @@ -27,7 +28,7 @@ def test_add_new_enum_value_render(connection: 'Connection'): database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.append('banned') + new_enum_variants.append("banned") target_schema = get_schema_with_enum_variants(new_enum_variants) @@ -36,21 +37,25 @@ def test_add_new_enum_value_render(connection: 'Connection'): template_args = {} autogenerate._render_migration_diffs(context, template_args) - assert (template_args["upgrades"] == - f"""# ### commands auto generated by Alembic - please adjust! ### + assert ( + template_args["upgrades"] + == f"""# ### commands auto generated by Alembic - please adjust! ### op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, new_enum_variants))}], [TableReference('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')], enum_values_to_rename=[]) - # ### end Alembic commands ###""") - assert (template_args["downgrades"] == - f"""# ### commands auto generated by Alembic - please adjust! ### + # ### end Alembic commands ###""" + ) + assert ( + template_args["downgrades"] + == f"""# ### commands auto generated by Alembic - please adjust! ### op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, old_enum_variants))}], [TableReference('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')], enum_values_to_rename=[]) - # ### end Alembic commands ###""") + # ### end Alembic commands ###""" + ) -def test_add_new_enum_value_diff_tuple(connection: 'Connection'): +def test_add_new_enum_value_diff_tuple(connection: "Connection"): """Check that enum variants are updated when new variant is added""" old_enum_variants = ["active", "passive"] @@ -58,7 +63,7 @@ def test_add_new_enum_value_diff_tuple(connection: 'Connection'): database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.append('banned') + new_enum_variants.append("banned") target_schema = get_schema_with_enum_variants(new_enum_variants) @@ -77,11 +82,11 @@ def test_add_new_enum_value_diff_tuple(connection: 'Connection'): SyncEnumValuesOp.operation_name, old_enum_variants, new_enum_variants, - [TableReference(USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)] + [TableReference(USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)], ) -def test_remove_enum_value_diff_tuple(connection: 'Connection'): +def test_remove_enum_value_diff_tuple(connection: "Connection"): """Check that enum variants are updated when new variant is removed""" old_enum_variants = ["active", "passive", "banned"] @@ -89,7 +94,7 @@ def test_remove_enum_value_diff_tuple(connection: 'Connection'): database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('banned') + new_enum_variants.remove("banned") target_schema = get_schema_with_enum_variants(new_enum_variants) @@ -104,7 +109,12 @@ def test_remove_enum_value_diff_tuple(connection: 'Connection'): assert len(diffs) == 1 change_variants_diff_tuple = diffs[0] - operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple + ( + operation_name, + old_values, + new_values, + affected_columns, + ) = change_variants_diff_tuple assert operation_name == SyncEnumValuesOp.operation_name assert old_values == old_enum_variants @@ -114,7 +124,7 @@ def test_remove_enum_value_diff_tuple(connection: 'Connection'): ] -def test_rename_enum_value_diff_tuple(connection: 'Connection'): +def test_rename_enum_value_diff_tuple(connection: "Connection"): """Check that enum variants are updated when a variant is renamed""" old_enum_variants = ["active", "passive", "banned"] @@ -122,8 +132,8 @@ def test_rename_enum_value_diff_tuple(connection: 'Connection'): database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('banned') - new_enum_variants.append('inactive') + new_enum_variants.remove("banned") + new_enum_variants.append("inactive") target_schema = get_schema_with_enum_variants(new_enum_variants) @@ -138,7 +148,12 @@ def test_rename_enum_value_diff_tuple(connection: 'Connection'): assert len(diffs) == 1 change_variants_diff_tuple = diffs[0] - operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple + ( + operation_name, + old_values, + new_values, + affected_columns, + ) = change_variants_diff_tuple assert operation_name == SyncEnumValuesOp.operation_name assert old_values == old_enum_variants @@ -146,5 +161,3 @@ def test_rename_enum_value_diff_tuple(connection: 'Connection'): assert affected_columns == [ TableReference(USER_TABLE_NAME, USER_STATUS_COLUMN_NAME) ] - - diff --git a/tests/sync_enum_values/test_run_in_different_schema.py b/tests/sync_enum_values/test_run_in_different_schema.py index 7ab26f8..d20078b 100644 --- a/tests/sync_enum_values/test_run_in_different_schema.py +++ b/tests/sync_enum_values/test_run_in_different_schema.py @@ -35,7 +35,7 @@ class TableWithExplicitEnumSchema(Base): ) -def test_run_in_different_schema(connection: 'Connection'): +def test_run_in_different_schema(connection: "Connection"): """https://github.com/Pogchamp-company/alembic-postgresql-enum/issues/34""" old_enum_variants = list(map(lambda item: item.name, _TestStatus)) @@ -43,16 +43,18 @@ def test_run_in_different_schema(connection: 'Connection'): database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.append('WAITING_FOR_APPROVAL') + new_enum_variants.append("WAITING_FOR_APPROVAL") mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(Base.metadata.schema, 'test_status', new_enum_variants, - [(TableWithExplicitEnumSchema.__tablename__, 'status')]) + ops.sync_enum_values( + Base.metadata.schema, + "test_status", + new_enum_variants, + [(TableWithExplicitEnumSchema.__tablename__, "status")], + ) defined = get_defined_enums(connection, Base.metadata.schema) - assert defined == { - 'test_status': tuple(new_enum_variants) - } + assert defined == {"test_status": tuple(new_enum_variants)} diff --git a/tests/sync_enum_values/test_run_sync_enum_values.py b/tests/sync_enum_values/test_run_sync_enum_values.py index 0c258c6..1ff93c1 100644 --- a/tests/sync_enum_values/test_run_sync_enum_values.py +++ b/tests/sync_enum_values/test_run_sync_enum_values.py @@ -7,183 +7,233 @@ from alembic_postgresql_enum.get_enum_data import ColumnType, get_defined_enums from alembic_postgresql_enum.sql_commands.column_default import get_column_default -from tests.schemas import (get_schema_with_enum_variants, DEFAULT_SCHEMA, - USER_STATUS_ENUM_NAME, USER_STATUS_COLUMN_NAME, - USER_TABLE_NAME, get_schema_with_enum_in_array_variants, CAR_TABLE_NAME, - CAR_COLORS_COLUMN_NAME, CAR_COLORS_ENUM_NAME) - - -def test_sync_enum_values_with_new_value(connection: 'Connection'): +from tests.schemas import ( + get_schema_with_enum_variants, + DEFAULT_SCHEMA, + USER_STATUS_ENUM_NAME, + USER_STATUS_COLUMN_NAME, + USER_TABLE_NAME, + get_schema_with_enum_in_array_variants, + CAR_TABLE_NAME, + CAR_COLORS_COLUMN_NAME, + CAR_COLORS_ENUM_NAME, +) + + +def test_sync_enum_values_with_new_value(connection: "Connection"): old_enum_variants = ["active", "passive"] database_schema = get_schema_with_enum_variants(old_enum_variants) database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.append('banned') + new_enum_variants.append("banned") mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(DEFAULT_SCHEMA, USER_STATUS_ENUM_NAME, new_enum_variants, - ((USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),)) + ops.sync_enum_values( + DEFAULT_SCHEMA, + USER_STATUS_ENUM_NAME, + new_enum_variants, + ((USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),), + ) defined = get_defined_enums(connection, DEFAULT_SCHEMA) - assert defined == { - USER_STATUS_ENUM_NAME: tuple(new_enum_variants) - } + assert defined == {USER_STATUS_ENUM_NAME: tuple(new_enum_variants)} -def test_sync_enum_values_with_renamed_value(connection: 'Connection'): +def test_sync_enum_values_with_renamed_value(connection: "Connection"): old_enum_variants = ["active", "passive"] database_schema = get_schema_with_enum_variants(old_enum_variants) database_schema.create_all(connection) - connection.execute(sqlalchemy.text(f''' + connection.execute( + sqlalchemy.text( + f""" INSERT INTO {USER_TABLE_NAME} ({USER_STATUS_COLUMN_NAME}) VALUES ('active'), ('passive') - ''')) + """ + ) + ) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('passive') - new_enum_variants.append('inactive') + new_enum_variants.remove("passive") + new_enum_variants.append("inactive") mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(DEFAULT_SCHEMA, USER_STATUS_ENUM_NAME, new_enum_variants, - ((USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),), - enum_values_to_rename=[ - ('passive', 'inactive') - ]) + ops.sync_enum_values( + DEFAULT_SCHEMA, + USER_STATUS_ENUM_NAME, + new_enum_variants, + ((USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),), + enum_values_to_rename=[("passive", "inactive")], + ) defined = get_defined_enums(connection, DEFAULT_SCHEMA) - assert defined == { - USER_STATUS_ENUM_NAME: tuple(new_enum_variants) - } + assert defined == {USER_STATUS_ENUM_NAME: tuple(new_enum_variants)} - users_entries = connection.execute(sqlalchemy.text(f''' + users_entries = ( + connection.execute( + sqlalchemy.text( + f""" SELECT {USER_STATUS_COLUMN_NAME} FROM {USER_TABLE_NAME} - ''')).scalars().all() + """ + ) + ) + .scalars() + .all() + ) - assert users_entries == ['active', 'inactive'] + assert users_entries == ["active", "inactive"] -def test_sync_enum_values_with_server_default(connection: 'Connection'): +def test_sync_enum_values_with_server_default(connection: "Connection"): old_enum_variants = ["active", "passive"] database_schema = MetaData() Table( - 'orders', + "orders", database_schema, Column("id", Integer, primary_key=True), - Column('status', sqlalchemy.Enum(*old_enum_variants, name='order_status'), - server_default=old_enum_variants[0]) + Column( + "status", + sqlalchemy.Enum(*old_enum_variants, name="order_status"), + server_default=old_enum_variants[0], + ), ) database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants.append('banned') + new_enum_variants.append("banned") mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(DEFAULT_SCHEMA, 'order_status', new_enum_variants, - (('orders', 'status'),)) + ops.sync_enum_values( + DEFAULT_SCHEMA, "order_status", new_enum_variants, (("orders", "status"),) + ) defined = get_defined_enums(connection, DEFAULT_SCHEMA) - assert defined == { - 'order_status': tuple(new_enum_variants) - } + assert defined == {"order_status": tuple(new_enum_variants)} -def test_sync_enum_values_with_server_default_renamed(connection: 'Connection'): +def test_sync_enum_values_with_server_default_renamed(connection: "Connection"): old_enum_variants = ["active", "passive"] database_schema = MetaData() Table( - 'orders', + "orders", database_schema, Column("id", Integer, primary_key=True), - Column('status', sqlalchemy.Enum(*old_enum_variants, name='order_status'), - server_default=old_enum_variants[1]) + Column( + "status", + sqlalchemy.Enum(*old_enum_variants, name="order_status"), + server_default=old_enum_variants[1], + ), ) database_schema.create_all(connection) new_enum_variants = old_enum_variants.copy() - new_enum_variants[1] = 'inactive' + new_enum_variants[1] = "inactive" mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(DEFAULT_SCHEMA, 'order_status', new_enum_variants, - (('orders', 'status'),), - enum_values_to_rename=[('passive', 'inactive')]) + ops.sync_enum_values( + DEFAULT_SCHEMA, + "order_status", + new_enum_variants, + (("orders", "status"),), + enum_values_to_rename=[("passive", "inactive")], + ) defined = get_defined_enums(connection, DEFAULT_SCHEMA) - order_status_default = get_column_default(connection, DEFAULT_SCHEMA, 'orders', 'status') + order_status_default = get_column_default( + connection, DEFAULT_SCHEMA, "orders", "status" + ) assert order_status_default == "'inactive'::order_status" - assert defined == { - 'order_status': tuple(new_enum_variants) - } + assert defined == {"order_status": tuple(new_enum_variants)} -def test_sync_enum_values_raise_custom_exception(connection: 'Connection'): +def test_sync_enum_values_raise_custom_exception(connection: "Connection"): old_enum_variants = ["active", "passive", "banned"] database_schema = get_schema_with_enum_variants(old_enum_variants) database_schema.create_all(connection) - connection.execute(sqlalchemy.text(f''' + connection.execute( + sqlalchemy.text( + f""" INSERT INTO {USER_TABLE_NAME} ({USER_STATUS_COLUMN_NAME}) VALUES ('active'), ('passive'), ('banned') - ''')) + """ + ) + ) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('banned') + new_enum_variants.remove("banned") mc = MigrationContext.configure(connection) ops = Operations(mc) with pytest.raises(ValueError): - ops.sync_enum_values(DEFAULT_SCHEMA, USER_STATUS_ENUM_NAME, new_enum_variants, - ((USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),), - enum_values_to_rename=[]) + ops.sync_enum_values( + DEFAULT_SCHEMA, + USER_STATUS_ENUM_NAME, + new_enum_variants, + ((USER_TABLE_NAME, USER_STATUS_COLUMN_NAME),), + enum_values_to_rename=[], + ) -def test_sync_enum_values_with_renamed_value_with_array(connection: 'Connection'): - old_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'violet', 'other'] +def test_sync_enum_values_with_renamed_value_with_array(connection: "Connection"): + old_enum_variants = ["black", "white", "red", "green", "blue", "violet", "other"] database_schema = get_schema_with_enum_in_array_variants(old_enum_variants) database_schema.create_all(connection) - connection.execute(sqlalchemy.text(f''' + connection.execute( + sqlalchemy.text( + f""" INSERT INTO {CAR_TABLE_NAME} ({CAR_COLORS_COLUMN_NAME}) VALUES ('{{"black"}}'), ('{{"white", "violet"}}') - ''')) + """ + ) + ) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('violet') - new_enum_variants.append('purple') + new_enum_variants.remove("violet") + new_enum_variants.append("purple") mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(DEFAULT_SCHEMA, CAR_COLORS_ENUM_NAME, new_enum_variants, - ((CAR_TABLE_NAME, CAR_COLORS_COLUMN_NAME, ColumnType.ARRAY),), - enum_values_to_rename=[ - ('violet', 'purple') - ]) + ops.sync_enum_values( + DEFAULT_SCHEMA, + CAR_COLORS_ENUM_NAME, + new_enum_variants, + ((CAR_TABLE_NAME, CAR_COLORS_COLUMN_NAME, ColumnType.ARRAY),), + enum_values_to_rename=[("violet", "purple")], + ) defined = get_defined_enums(connection, DEFAULT_SCHEMA) - assert defined == { - CAR_COLORS_ENUM_NAME: tuple(new_enum_variants) - } + assert defined == {CAR_COLORS_ENUM_NAME: tuple(new_enum_variants)} - users_entries = connection.execute(sqlalchemy.text(f''' + users_entries = ( + connection.execute( + sqlalchemy.text( + f""" SELECT {CAR_COLORS_COLUMN_NAME} FROM {CAR_TABLE_NAME} - ''')).scalars().all() + """ + ) + ) + .scalars() + .all() + ) - assert users_entries == ['{black}', '{white,purple}'] + assert users_entries == ["{black}", "{white,purple}"] diff --git a/tests/test_alter_column/test_text_column.py b/tests/test_alter_column/test_text_column.py index 02bd932..d213963 100644 --- a/tests/test_alter_column/test_text_column.py +++ b/tests/test_alter_column/test_text_column.py @@ -11,30 +11,33 @@ class NewEnum(Enum): - A = 'a' - B = 'b' - C = 'c' + A = "a" + B = "b" + C = "c" -def test_text_column(connection: 'Connection'): +def test_text_column(connection: "Connection"): database_schema = MetaData() - a_table = Table('a', database_schema, - Column('value', TEXT) - ) + 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}, - ])) + 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)) - ) + Table("a", target_schema, Column("value", postgresql.ENUM(NewEnum))) - compare_and_run(connection, target_schema, expected_upgrade=f""" + 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', @@ -43,7 +46,8 @@ def test_text_column(connection: 'Connection'): existing_nullable=True, postgresql_using='value::newenum') # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=f""" # ### commands auto generated by Alembic - please adjust! ### op.alter_column('a', 'value', existing_type=postgresql.ENUM('A', 'B', 'C', name='newenum'), @@ -51,4 +55,5 @@ def test_text_column(connection: 'Connection'): existing_nullable=True) sa.Enum('A', 'B', 'C', name='newenum').drop(op.get_bind()) # ### end Alembic commands ### - """) + """, + ) diff --git a/tests/test_constraints/test_unique.py b/tests/test_constraints/test_unique.py index 496b194..1e4fa9a 100644 --- a/tests/test_constraints/test_unique.py +++ b/tests/test_constraints/test_unique.py @@ -8,7 +8,7 @@ from tests.schemas import USER_TABLE_NAME, DEFAULT_SCHEMA -def test_sync_enum_values_unique_column_with_equals(connection: 'Connection'): +def test_sync_enum_values_unique_column_with_equals(connection: "Connection"): old_enum_variants = ["admin", "writer"] database_schema = MetaData() @@ -17,31 +17,43 @@ def test_sync_enum_values_unique_column_with_equals(connection: 'Connection'): USER_TABLE_NAME, database_schema, Column("id", Integer, primary_key=True), - Column('role', postgresql.ENUM(*old_enum_variants, name='user_role'), - nullable=True), - Index('admin_unique', 'role', unique=True, postgresql_where=sqlalchemy.text(f"role = 'admin'::user_role")) + Column( + "role", postgresql.ENUM(*old_enum_variants, name="user_role"), nullable=True + ), + Index( + "admin_unique", + "role", + unique=True, + postgresql_where=sqlalchemy.text(f"role = 'admin'::user_role"), + ), ) database_schema.create_all(connection) - connection.execute(sqlalchemy.text(f''' + connection.execute( + sqlalchemy.text( + f""" INSERT INTO {USER_TABLE_NAME} (role) VALUES ('admin'), ('writer'), (null) - ''')) + """ + ) + ) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('admin') - new_enum_variants.insert(0, 'administrator') + new_enum_variants.remove("admin") + new_enum_variants.insert(0, "administrator") mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(DEFAULT_SCHEMA, 'user_role', new_enum_variants, - [(USER_TABLE_NAME, 'role')], - enum_values_to_rename=[ - ('admin', 'administrator') - ]) + ops.sync_enum_values( + DEFAULT_SCHEMA, + "user_role", + new_enum_variants, + [(USER_TABLE_NAME, "role")], + enum_values_to_rename=[("admin", "administrator")], + ) -def test_sync_enum_values_unique_column_with_not_equals(connection: 'Connection'): +def test_sync_enum_values_unique_column_with_not_equals(connection: "Connection"): old_enum_variants = ["admin", "writer"] database_schema = MetaData() @@ -50,25 +62,37 @@ def test_sync_enum_values_unique_column_with_not_equals(connection: 'Connection' USER_TABLE_NAME, database_schema, Column("id", Integer, primary_key=True), - Column('role', postgresql.ENUM(*old_enum_variants, name='user_role'), - nullable=True), - Index('admin_unique', 'role', unique=True, postgresql_where=sqlalchemy.text(f"role <> 'admin'::user_role")) + Column( + "role", postgresql.ENUM(*old_enum_variants, name="user_role"), nullable=True + ), + Index( + "admin_unique", + "role", + unique=True, + postgresql_where=sqlalchemy.text(f"role <> 'admin'::user_role"), + ), ) database_schema.create_all(connection) - connection.execute(sqlalchemy.text(f''' + connection.execute( + sqlalchemy.text( + f""" INSERT INTO {USER_TABLE_NAME} (role) VALUES ('admin'), ('writer'), (null) - ''')) + """ + ) + ) new_enum_variants = old_enum_variants.copy() - new_enum_variants.remove('admin') - new_enum_variants.insert(0, 'administrator') + new_enum_variants.remove("admin") + new_enum_variants.insert(0, "administrator") mc = MigrationContext.configure(connection) ops = Operations(mc) - ops.sync_enum_values(DEFAULT_SCHEMA, 'user_role', new_enum_variants, - [(USER_TABLE_NAME, 'role')], - enum_values_to_rename=[ - ('admin', 'administrator') - ]) + ops.sync_enum_values( + DEFAULT_SCHEMA, + "user_role", + new_enum_variants, + [(USER_TABLE_NAME, "role")], + enum_values_to_rename=[("admin", "administrator")], + ) diff --git a/tests/test_enum_creation/test_add_column.py b/tests/test_enum_creation/test_add_column.py index 6b60002..dcb60f2 100644 --- a/tests/test_enum_creation/test_add_column.py +++ b/tests/test_enum_creation/test_add_column.py @@ -7,9 +7,17 @@ from alembic_postgresql_enum.operations import CreateEnumOp from tests.base.render_and_run import compare_and_run -from tests.schemas import get_schema_with_enum_variants, USER_TABLE_NAME, USER_STATUS_ENUM_NAME, \ - get_schema_without_enum, USER_STATUS_COLUMN_NAME, DEFAULT_SCHEMA, CAR_COLORS_ENUM_NAME, \ - get_schema_with_enum_in_array_variants, get_car_schema_without_enum +from tests.schemas import ( + get_schema_with_enum_variants, + USER_TABLE_NAME, + USER_STATUS_ENUM_NAME, + get_schema_without_enum, + USER_STATUS_COLUMN_NAME, + DEFAULT_SCHEMA, + CAR_COLORS_ENUM_NAME, + get_schema_with_enum_in_array_variants, + get_car_schema_without_enum, +) from tests.utils.migration_context import create_migration_context if TYPE_CHECKING: @@ -17,7 +25,7 @@ from sqlalchemy import MetaData, Table, Column, Integer -def test_create_enum_before_add_column(connection: 'Connection'): +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() database_schema.create_all(connection) @@ -26,20 +34,25 @@ def test_create_enum_before_add_column(connection: 'Connection'): target_schema = get_schema_with_enum_variants(new_enum_variants) - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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)) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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()) # ### end Alembic commands ### - """) + """, + ) -def test_create_enum_diff_tuple(connection: 'Connection'): +def test_create_enum_diff_tuple(connection: "Connection"): """Check that library correctly creates enum before its use inside add_column""" database_schema = get_schema_without_enum() database_schema.create_all(connection) @@ -63,17 +76,17 @@ def test_create_enum_diff_tuple(connection: 'Connection'): CreateEnumOp.operation_name, USER_STATUS_ENUM_NAME, DEFAULT_SCHEMA, - tuple(new_enum_variants) + tuple(new_enum_variants), ) - assert add_column_tuple[0] == 'add_column' + assert add_column_tuple[0] == "add_column" -def test_create_enum_diff_tuple_with_array(connection: 'Connection'): +def test_create_enum_diff_tuple_with_array(connection: "Connection"): """Check that library correctly creates enum before its use inside add_column""" database_schema = get_car_schema_without_enum() database_schema.create_all(connection) - new_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'other'] + new_enum_variants = ["black", "white", "red", "green", "blue", "other"] target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) @@ -92,12 +105,12 @@ def test_create_enum_diff_tuple_with_array(connection: 'Connection'): CreateEnumOp.operation_name, CAR_COLORS_ENUM_NAME, DEFAULT_SCHEMA, - tuple(new_enum_variants) + tuple(new_enum_variants), ) - assert add_column_tuple[0] == 'add_column' + assert add_column_tuple[0] == "add_column" -def test_with_non_native_enum(connection: 'Connection'): +def test_with_non_native_enum(connection: "Connection"): """Check that library ignores sa.Enum that are not native""" database_schema = get_schema_without_enum() database_schema.create_all(connection) @@ -110,22 +123,31 @@ def test_with_non_native_enum(connection: 'Connection'): 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)) + Column( + USER_STATUS_COLUMN_NAME, + sqlalchemy.Enum( + *new_enum_variants, name=USER_STATUS_ENUM_NAME, native_enum=False + ), + ), ) - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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)) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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'): +def test_create_enum_before_add_column_metadata_list(connection: "Connection"): """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) @@ -134,14 +156,19 @@ def test_create_enum_before_add_column_metadata_list(connection: 'Connection'): target_schema = get_schema_with_enum_variants(new_enum_variants) - compare_and_run(connection, [target_schema], expected_upgrade=f""" + compare_and_run( + connection, + [target_schema], + expected_upgrade=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)) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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()) # ### end Alembic commands ### - """) + """, + ) diff --git a/tests/test_enum_creation/test_create_schema.py b/tests/test_enum_creation/test_create_schema.py index 30d35a9..8144340 100644 --- a/tests/test_enum_creation/test_create_schema.py +++ b/tests/test_enum_creation/test_create_schema.py @@ -4,8 +4,11 @@ from alembic import autogenerate from sqlalchemy.dialects import postgresql -from tests.schemas import USER_TABLE_NAME, USER_STATUS_ENUM_NAME, \ - USER_STATUS_COLUMN_NAME +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: @@ -13,10 +16,10 @@ from sqlalchemy import Table, Column, Integer, MetaData -def test_create_enum_on_create_table_inside_new_schema(connection: 'Connection'): +def test_create_enum_on_create_table_inside_new_schema(connection: "Connection"): """Check that library correctly creates enum before its use inside create_table inside new schema""" new_enum_variants = ["active", "passive"] - non_existing_schema = 'non_existing_schema' + non_existing_schema = "non_existing_schema" target_schema = MetaData(schema=non_existing_schema) @@ -24,10 +27,12 @@ def test_create_enum_on_create_table_inside_new_schema(connection: 'Connection') 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 - )), + Column( + USER_STATUS_COLUMN_NAME, + postgresql.ENUM( + *new_enum_variants, name=USER_STATUS_ENUM_NAME, metadata=target_schema + ), + ), ) expected_upgrade = f""" @@ -54,11 +59,15 @@ def test_create_enum_on_create_table_inside_new_schema(connection: 'Connection') # 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"]) + 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 ') + 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}' + 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 59fa130..c1f0fae 100644 --- a/tests/test_enum_creation/test_create_table.py +++ b/tests/test_enum_creation/test_create_table.py @@ -4,22 +4,34 @@ from sqlalchemy.dialects import postgresql from tests.base.render_and_run import compare_and_run -from tests.schemas import get_schema_with_enum_variants, USER_TABLE_NAME, USER_STATUS_ENUM_NAME, \ - get_car_schema_without_enum, get_schema_with_enum_in_array_variants, CAR_TABLE_NAME, CAR_COLORS_COLUMN_NAME, \ - CAR_COLORS_ENUM_NAME, USER_STATUS_COLUMN_NAME, ANOTHER_SCHEMA_NAME +from tests.schemas import ( + get_schema_with_enum_variants, + USER_TABLE_NAME, + USER_STATUS_ENUM_NAME, + get_car_schema_without_enum, + get_schema_with_enum_in_array_variants, + CAR_TABLE_NAME, + CAR_COLORS_COLUMN_NAME, + CAR_COLORS_ENUM_NAME, + USER_STATUS_COLUMN_NAME, + 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'): +def test_create_enum_on_create_table(connection: "Connection"): """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) - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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}', @@ -28,37 +40,44 @@ def test_create_enum_on_create_table(connection: 'Connection'): sa.PrimaryKeyConstraint('id') ) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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()) # ### end Alembic commands ### - """) + """, + ) -def test_create_enum_on_create_table_with_array(connection: 'Connection'): +def test_create_enum_on_create_table_with_array(connection: "Connection"): """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'] + new_enum_variants = ["black", "white", "red", "green", "blue", "other"] target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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}')), nullable=True)) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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()) # ### end Alembic commands ### - """) + """, + ) -def test_create_enum_on_create_table_with_sa_enum(connection: 'Connection'): +def test_create_enum_on_create_table_with_sa_enum(connection: "Connection"): """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 """ @@ -70,10 +89,16 @@ def test_create_enum_on_create_table_with_sa_enum(connection: 'Connection'): 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)) + Column( + USER_STATUS_COLUMN_NAME, + sqlalchemy.Enum(*new_enum_variants, name=USER_STATUS_ENUM_NAME), + ), ) - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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}', @@ -82,15 +107,17 @@ def test_create_enum_on_create_table_with_sa_enum(connection: 'Connection'): sa.PrimaryKeyConstraint('id') ) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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()) # ### end Alembic commands ### - """) + """, + ) -def test_create_enum_on_create_table_with_another_schema(connection: 'Connection'): +def test_create_enum_on_create_table_with_another_schema(connection: "Connection"): """Check that library correctly creates enum before its use inside create_table inside another schema""" new_enum_variants = ["active", "passive"] @@ -100,13 +127,21 @@ def test_create_enum_on_create_table_with_another_schema(connection: 'Connection 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)), - schema=ANOTHER_SCHEMA_NAME + Column( + USER_STATUS_COLUMN_NAME, + postgresql.ENUM( + *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""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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()) op.create_table('{USER_TABLE_NAME}', @@ -116,15 +151,19 @@ def test_create_enum_on_create_table_with_another_schema(connection: 'Connection schema='{ANOTHER_SCHEMA_NAME}' ) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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()) # ### end Alembic commands ### - """) + """, + ) -def test_create_enum_on_create_table_with_another_schema_on_metadata(connection: 'Connection'): +def test_create_enum_on_create_table_with_another_schema_on_metadata( + connection: "Connection", +): """Check that library correctly creates enum before its use inside create_table inside another schema, specified on Metadata""" new_enum_variants = ["active", "passive"] @@ -134,13 +173,18 @@ def test_create_enum_on_create_table_with_another_schema_on_metadata(connection: 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 - )), + 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""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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()) op.create_table('{USER_TABLE_NAME}', @@ -150,9 +194,11 @@ def test_create_enum_on_create_table_with_another_schema_on_metadata(connection: schema='{ANOTHER_SCHEMA_NAME}' ) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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()) # ### end Alembic commands ### - """) + """, + ) diff --git a/tests/test_enum_creation/test_drop_column.py b/tests/test_enum_creation/test_drop_column.py index 8c3910a..b8c07c7 100644 --- a/tests/test_enum_creation/test_drop_column.py +++ b/tests/test_enum_creation/test_drop_column.py @@ -5,11 +5,13 @@ if TYPE_CHECKING: from sqlalchemy import Connection -from tests.schemas import (get_schema_with_enum_in_array_variants, - CAR_TABLE_NAME, - CAR_COLORS_COLUMN_NAME, - CAR_COLORS_ENUM_NAME, get_car_schema_without_enum - ) +from tests.schemas import ( + get_schema_with_enum_in_array_variants, + CAR_TABLE_NAME, + CAR_COLORS_COLUMN_NAME, + CAR_COLORS_ENUM_NAME, + get_car_schema_without_enum, +) from typing import TYPE_CHECKING from alembic import autogenerate @@ -17,15 +19,21 @@ 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, USER_STATUS_ENUM_NAME, \ - get_schema_without_enum, USER_STATUS_COLUMN_NAME, DEFAULT_SCHEMA +from tests.schemas import ( + get_schema_with_enum_variants, + USER_TABLE_NAME, + USER_STATUS_ENUM_NAME, + get_schema_without_enum, + USER_STATUS_COLUMN_NAME, + DEFAULT_SCHEMA, +) from tests.utils.migration_context import create_migration_context if TYPE_CHECKING: from sqlalchemy import Connection -def test_delete_enum_after_drop_column(connection: 'Connection'): +def test_delete_enum_after_drop_column(connection: "Connection"): """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) @@ -33,41 +41,51 @@ def test_delete_enum_after_drop_column(connection: 'Connection'): target_schema = get_schema_without_enum() - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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()) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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)) # ### end Alembic commands ### - """) # For some reason alembic decided to add redundant autoincrement=False on downgrade + """, + ) # For some reason alembic decided to add redundant autoincrement=False on downgrade -def test_delete_enum_after_drop_column_with_array(connection: 'Connection'): +def test_delete_enum_after_drop_column_with_array(connection: "Connection"): """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'] + 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() - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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()) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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}')), autoincrement=False, nullable=True)) # ### end Alembic commands ### - """) + """, + ) -def test_delete_enum_diff_tuple(connection: 'Connection'): +def test_delete_enum_diff_tuple(connection: "Connection"): """Check that library correctly removes unused enum""" old_enum_variants = ["active", "passive"] database_schema = get_schema_with_enum_variants(old_enum_variants) @@ -86,18 +104,18 @@ def test_delete_enum_diff_tuple(connection: 'Connection'): assert len(diffs) == 2 remove_column_tuple, create_enum_tuple = diffs - assert remove_column_tuple[0] == 'remove_column' + assert remove_column_tuple[0] == "remove_column" assert create_enum_tuple == ( DropEnumOp.operation_name, USER_STATUS_ENUM_NAME, DEFAULT_SCHEMA, - tuple(old_enum_variants) + tuple(old_enum_variants), ) -def test_delete_enum_diff_tuple_with_array(connection: 'Connection'): +def test_delete_enum_diff_tuple_with_array(connection: "Connection"): """Check that library correctly removes unused enum""" - old_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'other'] + old_enum_variants = ["black", "white", "red", "green", "blue", "other"] database_schema = get_schema_with_enum_in_array_variants(old_enum_variants) database_schema.create_all(connection) @@ -114,10 +132,10 @@ def test_delete_enum_diff_tuple_with_array(connection: 'Connection'): assert len(diffs) == 2 remove_column_tuple, create_enum_tuple = diffs - assert remove_column_tuple[0] == 'remove_column' + assert remove_column_tuple[0] == "remove_column" assert create_enum_tuple == ( DropEnumOp.operation_name, CAR_COLORS_ENUM_NAME, DEFAULT_SCHEMA, - tuple(old_enum_variants) + tuple(old_enum_variants), ) diff --git a/tests/test_enum_creation/test_drop_table.py b/tests/test_enum_creation/test_drop_table.py index b77ac93..9af2b39 100644 --- a/tests/test_enum_creation/test_drop_table.py +++ b/tests/test_enum_creation/test_drop_table.py @@ -3,13 +3,17 @@ from sqlalchemy import MetaData from tests.base.render_and_run import compare_and_run -from tests.schemas import get_schema_with_enum_variants, USER_TABLE_NAME, USER_STATUS_ENUM_NAME +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'): +def test_drop_enum_after_drop_table(connection: "Connection"): """Check that library correctly drop enum after drop_table""" dropped_enum_variants = ["active", "passive"] database_schema = get_schema_with_enum_variants(dropped_enum_variants) @@ -17,12 +21,16 @@ def test_drop_enum_after_drop_table(connection: 'Connection'): target_schema = MetaData() - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=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()) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=f""" # ### commands auto generated by Alembic - please adjust! ### sa.Enum('active', 'passive', name='user_status').create(op.get_bind()) op.create_table('users', @@ -31,4 +39,5 @@ 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 cbab999..1231419 100644 --- a/tests/test_without_schema_changes/test_explicit_schema.py +++ b/tests/test_without_schema_changes/test_explicit_schema.py @@ -2,7 +2,11 @@ import enum from typing import TYPE_CHECKING -from alembic_postgresql_enum.get_enum_data import TableReference, get_defined_enums, get_declared_enums +from alembic_postgresql_enum.get_enum_data import ( + TableReference, + get_defined_enums, + get_declared_enums, +) from tests.base.render_and_run import compare_and_run from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA @@ -36,42 +40,51 @@ class TableWithExplicitEnumSchema(Base): ) -def test_get_defined_enums(connection: 'Connection'): +def test_get_defined_enums(connection: "Connection"): database_schema = my_metadata database_schema.create_all(connection) function_result = get_defined_enums(connection, ANOTHER_SCHEMA_NAME) assert function_result == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } -def test_get_declared_enums(connection: 'Connection'): +def test_get_declared_enums(connection: "Connection"): declared_schema = my_metadata - function_result = get_declared_enums(declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection + ) assert function_result.enum_values == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } assert function_result.enum_table_references == { - 'test_status': frozenset([TableReference(TableWithExplicitEnumSchema.__tablename__, 'status')]) + "test_status": frozenset( + [TableReference(TableWithExplicitEnumSchema.__tablename__, "status")] + ) } -def test_compare_and_run(connection: 'Connection'): +def test_compare_and_run(connection: "Connection"): database_schema = my_metadata database_schema.create_all(connection) target_schema = my_metadata - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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 c54f98e..f7e9af5 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 @@ -2,7 +2,11 @@ import enum from typing import TYPE_CHECKING -from alembic_postgresql_enum.get_enum_data import TableReference, get_defined_enums, get_declared_enums +from alembic_postgresql_enum.get_enum_data import ( + TableReference, + get_defined_enums, + get_declared_enums, +) from tests.base.render_and_run import compare_and_run from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA @@ -36,47 +40,55 @@ class TableWithExplicitEnumSchema(Base): ) -def test_get_defined_enums_metadata(connection: 'Connection'): +def test_get_defined_enums_metadata(connection: "Connection"): database_schema = my_metadata database_schema.create_all(connection) # Check that enum is not created inside `another` schema function_result = get_defined_enums(connection, ANOTHER_SCHEMA_NAME) - assert function_result == { - } + assert function_result == {} # Check that enum is created inside `public` schema function_result = get_defined_enums(connection, DEFAULT_SCHEMA) assert function_result == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } -def test_get_declared_enums(connection: 'Connection'): +def test_get_declared_enums(connection: "Connection"): declared_schema = my_metadata # Check that enum is not created inside `another` schema - function_result = get_declared_enums(declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection + ) assert function_result.enum_values == {} assert function_result.enum_table_references == {} # Check that enum is created inside `public` schema - function_result = get_declared_enums(declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection + ) assert function_result.enum_values == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } assert function_result.enum_table_references == { - 'test_status': frozenset([TableReference(TableWithExplicitEnumSchema.__tablename__, 'status')]) + "test_status": frozenset( + [TableReference(TableWithExplicitEnumSchema.__tablename__, "status")] + ) } -def test_compare_and_run_create_table(connection: 'Connection'): +def test_compare_and_run_create_table(connection: "Connection"): target_schema = my_metadata - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=f""" # ### commands auto generated by Alembic - please adjust! ### sa.Enum('PENDING', 'SUCCESS', 'FAILED', name='test_status').create(op.get_bind()) op.create_table('test', @@ -86,26 +98,33 @@ def test_compare_and_run_create_table(connection: 'Connection'): schema='another' ) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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'): +def test_compare_and_run(connection: "Connection"): database_schema = my_metadata database_schema.create_all(connection) target_schema = my_metadata - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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 5b6959a..0c06607 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 @@ -5,7 +5,11 @@ import pytest import sqlalchemy -from alembic_postgresql_enum.get_enum_data import TableReference, get_defined_enums, get_declared_enums +from alembic_postgresql_enum.get_enum_data import ( + TableReference, + get_defined_enums, + get_declared_enums, +) from tests.base.render_and_run import compare_and_run from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA @@ -43,47 +47,60 @@ class TableWithExplicitEnumSchema(Base): return my_metadata -@pytest.mark.skipif(sqlalchemy.__version__.startswith('1.'), reason="Table are made in 2.0 style") -def test_get_defined_enums_metadata(connection: 'Connection'): +@pytest.mark.skipif( + sqlalchemy.__version__.startswith("1."), reason="Table are made in 2.0 style" +) +def test_get_defined_enums_metadata(connection: "Connection"): my_metadata = get_my_metadata() database_schema = my_metadata database_schema.create_all(connection) - function_result = get_defined_enums(connection, 'public') + function_result = get_defined_enums(connection, "public") assert function_result == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } -@pytest.mark.skipif(sqlalchemy.__version__.startswith('1.'), reason="Table are made in 2.0 style") -def test_get_declared_enums(connection: 'Connection'): +@pytest.mark.skipif( + sqlalchemy.__version__.startswith("1."), reason="Table are made in 2.0 style" +) +def test_get_declared_enums(connection: "Connection"): my_metadata = get_my_metadata() declared_schema = my_metadata # Check that enum is not created inside `another` schema - function_result = get_declared_enums(declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection + ) assert function_result.enum_values == {} assert function_result.enum_table_references == {} # Check that enum is created inside `public` schema - function_result = get_declared_enums(declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection + ) assert function_result.enum_values == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } assert function_result.enum_table_references == { - 'test_status': frozenset([TableReference('test', 'status')]) + "test_status": frozenset([TableReference("test", "status")]) } -@pytest.mark.skipif(sqlalchemy.__version__.startswith('1.'), reason="Table are made in 2.0 style") -def test_compare_and_run_create_table(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 - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=f""" # ### commands auto generated by Alembic - please adjust! ### sa.Enum('PENDING', 'SUCCESS', 'FAILED', name='test_status').create(op.get_bind()) op.create_table('test', @@ -93,28 +110,37 @@ def test_compare_and_run_create_table(connection: 'Connection'): schema='another' ) # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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'): +@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) target_schema = my_metadata - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=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 71be3c3..8e2cbf5 100644 --- a/tests/test_without_schema_changes/test_inherit_schema_true.py +++ b/tests/test_without_schema_changes/test_inherit_schema_true.py @@ -2,7 +2,11 @@ import enum from typing import TYPE_CHECKING -from alembic_postgresql_enum.get_enum_data import TableReference, get_defined_enums, get_declared_enums +from alembic_postgresql_enum.get_enum_data import ( + TableReference, + get_defined_enums, + get_declared_enums, +) from tests.base.render_and_run import compare_and_run from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA @@ -36,42 +40,51 @@ class TableWithExplicitEnumSchema(Base): ) -def test_get_defined_enums(connection: 'Connection'): +def test_get_defined_enums(connection: "Connection"): database_schema = my_metadata database_schema.create_all(connection) function_result = get_defined_enums(connection, ANOTHER_SCHEMA_NAME) assert function_result == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } -def test_get_declared_enums(connection: 'Connection'): +def test_get_declared_enums(connection: "Connection"): declared_schema = my_metadata - function_result = get_declared_enums(declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection) + function_result = get_declared_enums( + declared_schema, ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, connection + ) assert function_result.enum_values == { - 'test_status': tuple(map(lambda item: item.value, _TestStatus)) + "test_status": tuple(map(lambda item: item.value, _TestStatus)) } assert function_result.enum_table_references == { - 'test_status': frozenset([TableReference(TableWithExplicitEnumSchema.__tablename__, 'status')]) + "test_status": frozenset( + [TableReference(TableWithExplicitEnumSchema.__tablename__, "status")] + ) } -def test_compare_and_run(connection: 'Connection'): +def test_compare_and_run(connection: "Connection"): database_schema = my_metadata database_schema.create_all(connection) target_schema = my_metadata - compare_and_run(connection, target_schema, expected_upgrade=f""" + compare_and_run( + connection, + target_schema, + expected_upgrade=f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """, expected_downgrade=f""" + """, + expected_downgrade=f""" # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### - """) + """, + ) diff --git a/tests/utils/migration_context.py b/tests/utils/migration_context.py index 651646a..8e29ad3 100644 --- a/tests/utils/migration_context.py +++ b/tests/utils/migration_context.py @@ -22,7 +22,9 @@ def default_migration_options(target_schema: Union[MetaData, List[MetaData]]) -> } -def create_migration_context(connection: 'Connection', target_schema: Union[MetaData, List[MetaData]]): +def create_migration_context( + connection: "Connection", target_schema: Union[MetaData, List[MetaData]] +): return MigrationContext.configure( connection=connection, opts=default_migration_options(target_schema),