def test_autogen_process_directives( self, get_version_branch_path, use_separate_migration_branches): use_separate_migration_branches.return_value = True get_version_branch_path.side_effect = lambda cfg, release, branch: ( "/foo/expand" if branch == 'expand' else "/foo/contract") migration_script = alembic_ops.MigrationScript( 'eced083f5df', # these directives will be split into separate # expand/contract scripts alembic_ops.UpgradeOps( ops=[ alembic_ops.CreateTableOp( 'organization', [ sa.Column('id', sa.Integer(), primary_key=True), sa.Column('name', sa.String(50), nullable=False) ] ), alembic_ops.ModifyTableOps( 'user', ops=[ alembic_ops.AddColumnOp( 'user', sa.Column('organization_id', sa.Integer()) ), alembic_ops.CreateForeignKeyOp( 'org_fk', 'user', 'organization', ['organization_id'], ['id'] ), alembic_ops.DropConstraintOp( 'user', 'uq_user_org' ), alembic_ops.DropColumnOp( 'user', 'organization_name' ) ] ) ] ), # these will be discarded alembic_ops.DowngradeOps( ops=[ alembic_ops.AddColumnOp( 'user', sa.Column( 'organization_name', sa.String(50), nullable=True) ), alembic_ops.CreateUniqueConstraintOp( 'uq_user_org', 'user', ['user_name', 'organization_name'] ), alembic_ops.ModifyTableOps( 'user', ops=[ alembic_ops.DropConstraintOp('org_fk', 'user'), alembic_ops.DropColumnOp('user', 'organization_id') ] ), alembic_ops.DropTableOp('organization') ] ), message='create the organization table and ' 'replace user.organization_name' ) directives = [migration_script] autogen.process_revision_directives( mock.Mock(), mock.Mock(), directives ) expand = directives[0] contract = directives[1] self.assertEqual("/foo/expand", expand.version_path) self.assertEqual("/foo/contract", contract.version_path) self.assertTrue(expand.downgrade_ops.is_empty()) self.assertTrue(contract.downgrade_ops.is_empty()) self.assertEqual( textwrap.dedent("""\ ### commands auto generated by Alembic - please adjust! ### op.create_table('organization', sa.Column('id', sa.Integer(), nullable=False), sa.Column('name', sa.String(length=50), nullable=False), sa.PrimaryKeyConstraint('id') ) op.add_column('user', """ """sa.Column('organization_id', sa.Integer(), nullable=True)) op.create_foreign_key('org_fk', 'user', """ """'organization', ['organization_id'], ['id']) ### end Alembic commands ###"""), alembic_ag_api.render_python_code(expand.upgrade_ops) ) self.assertEqual( textwrap.dedent("""\ ### commands auto generated by Alembic - please adjust! ### op.drop_constraint('user', 'uq_user_org', type_=None) op.drop_column('user', 'organization_name') ### end Alembic commands ###"""), alembic_ag_api.render_python_code(contract.upgrade_ops) )
def test_autogen_process_directives(self, get_version_branch_path): get_version_branch_path.side_effect = lambda cfg, release, branch: ( "/foo/expand" if branch == 'expand' else "/foo/contract") migration_script = alembic_ops.MigrationScript( 'eced083f5df', # these directives will be split into separate # expand/contract scripts alembic_ops.UpgradeOps(ops=[ alembic_ops.CreateTableOp('organization', [ sa.Column('id', sa.Integer(), primary_key=True), sa.Column('name', sa.String(50), nullable=False) ]), alembic_ops.ModifyTableOps( 'user', ops=[ alembic_ops.AddColumnOp( 'user', sa.Column('organization_id', sa.Integer())), alembic_ops.CreateForeignKeyOp( 'org_fk', 'user', 'organization', ['organization_id'], ['id']), alembic_ops.DropConstraintOp('user', 'uq_user_org'), alembic_ops.DropColumnOp('user', 'organization_name') ]) ]), # these will be discarded alembic_ops.DowngradeOps(ops=[ alembic_ops.AddColumnOp( 'user', sa.Column( 'organization_name', sa.String(50), nullable=True)), alembic_ops.CreateUniqueConstraintOp( 'uq_user_org', 'user', ['user_name', 'organization_name']), alembic_ops.ModifyTableOps( 'user', ops=[ alembic_ops.DropConstraintOp('org_fk', 'user'), alembic_ops.DropColumnOp('user', 'organization_id') ]), alembic_ops.DropTableOp('organization') ]), message='create the organization table and ' 'replace user.organization_name') directives = [migration_script] autogen.process_revision_directives(mock.Mock(), mock.Mock(), directives) expand = directives[0] contract = directives[1] self.assertEqual("/foo/expand", expand.version_path) self.assertEqual("/foo/contract", contract.version_path) self.assertTrue(expand.downgrade_ops.is_empty()) self.assertTrue(contract.downgrade_ops.is_empty()) self.assertEqual( textwrap.dedent( """\ ### commands auto generated by Alembic - please adjust! ### op.create_table('organization', sa.Column('id', sa.Integer(), nullable=False), sa.Column('name', sa.String(length=50), nullable=False), sa.PrimaryKeyConstraint('id') ) op.add_column('user', """ """sa.Column('organization_id', sa.Integer(), nullable=True)) op.create_foreign_key('org_fk', 'user', """ """'organization', ['organization_id'], ['id']) ### end Alembic commands ###"""), alembic_ag_api.render_python_code(expand.upgrade_ops)) self.assertEqual( textwrap.dedent("""\ ### commands auto generated by Alembic - please adjust! ### op.drop_constraint('user', 'uq_user_org', type_=None) op.drop_column('user', 'organization_name') ### end Alembic commands ###"""), alembic_ag_api.render_python_code(contract.upgrade_ops))