- add tests for batch autogenerate

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 3f2bc8b441a94547f8e5453dc6a66979f630b8f1..4539a9c77286af065f6f38ccb4983e49f999084b 100644 (file)
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -152,6 +152,7 @@ def _produce_migration_diffs(context, template_args,
 
 def _indent(text):
     text = re.compile(r'^', re.M).sub("    ", text).strip()
+    text = re.compile(r' +$', re.M).sub("", text)
     return text
 
 

diff --git a/tests/test_autogenerate.py b/tests/test_autogenerate.py
index 046d9cb2649f6027b3e3981e50bb3ff02d2b5ce9..9f6cc002706e3692dca49804968acf0c93491a7c 100644 (file)
--- a/tests/test_autogenerate.py
+++ b/tests/test_autogenerate.py
@@ -45,7 +45,6 @@ class AutogenTest(object):
 
     configure_opts = {}
 
-    @classmethod
     def setup_class(cls):
         staging_env()
         cls.bind = cls._get_bind()
@@ -53,36 +52,37 @@ class AutogenTest(object):
         cls.m1.create_all(cls.bind)
         cls.m2 = cls._get_model_schema()
 
-        conn = cls.bind.connect()
+    @classmethod
+    def teardown_class(cls):
+        cls.m1.drop_all(cls.bind)
+        clear_staging_env()
+
+    def setUp(self):
+        conn = self.bind.connect()
         ctx_opts = {
             'compare_type': True,
             'compare_server_default': True,
-            'target_metadata': cls.m2,
+            'target_metadata': self.m2,
             'upgrade_token': "upgrades",
             'downgrade_token': "downgrades",
             'alembic_module_prefix': 'op.',
             'sqlalchemy_module_prefix': 'sa.',
         }
-        if cls.configure_opts:
-            ctx_opts.update(cls.configure_opts)
-        cls.context = context = MigrationContext.configure(
+        if self.configure_opts:
+            ctx_opts.update(self.configure_opts)
+        self.context = context = MigrationContext.configure(
             connection=conn,
             opts=ctx_opts
         )
 
         connection = context.bind
-        cls.autogen_context = {
+        self.autogen_context = {
             'imports': set(),
             'connection': connection,
             'dialect': connection.dialect,
             'context': context
         }
 
-    @classmethod
-    def teardown_class(cls):
-        cls.m1.drop_all(cls.bind)
-        clear_staging_env()
-
 
 class AutogenFixtureTest(object):
 
@@ -450,6 +450,34 @@ class AutogenerateDiffTest(ModelOne, AutogenTest, TestBase):
     pass
     ### end Alembic commands ###""")
 
+    def test_render_nothing_batch(self):
+        context = MigrationContext.configure(
+            connection=self.bind.connect(),
+            opts={
+                'compare_type': True,
+                'compare_server_default': True,
+                'target_metadata': self.m1,
+                'upgrade_token': "upgrades",
+                'downgrade_token': "downgrades",
+                'alembic_module_prefix': 'op.',
+                'sqlalchemy_module_prefix': 'sa.',
+                'render_as_batch': True
+            }
+        )
+        template_args = {}
+        autogenerate._produce_migration_diffs(
+            context, template_args, set(),
+            include_symbol=lambda name, schema: False
+        )
+        eq_(re.sub(r"u'", "'", template_args['upgrades']),
+            """### commands auto generated by Alembic - please adjust! ###
+    pass
+    ### end Alembic commands ###""")
+        eq_(re.sub(r"u'", "'", template_args['downgrades']),
+            """### commands auto generated by Alembic - please adjust! ###
+    pass
+    ### end Alembic commands ###""")
+
     def test_render_diffs_standard(self):
         """test a full render including indentation"""
 
@@ -512,6 +540,79 @@ nullable=True))
     op.drop_table('item')
     ### end Alembic commands ###""")
 
+    def test_render_diffs_batch(self):
+        """test a full render in batch mode including indentation"""
+
+        template_args = {}
+        self.context.opts['render_as_batch'] = True
+        autogenerate._produce_migration_diffs(
+            self.context, template_args, set())
+
+        eq_(re.sub(r"u'", "'", template_args['upgrades']),
+            """### commands auto generated by Alembic - please adjust! ###
+    op.create_table('item',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('description', sa.String(length=100), nullable=True),
+    sa.Column('order_id', sa.Integer(), nullable=True),
+    sa.CheckConstraint('len(description) > 5'),
+    sa.ForeignKeyConstraint(['order_id'], ['order.order_id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.drop_table('extra')
+    with op.batch_alter_table('address', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('street', sa.String(length=50), nullable=True))
+
+    with op.batch_alter_table('order', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('user_id', sa.Integer(), nullable=True))
+        batch_op.alter_column('amount',
+               existing_type=sa.NUMERIC(precision=8, scale=2),
+               type_=sa.Numeric(precision=10, scale=2),
+               nullable=True,
+               existing_server_default=sa.text('0'))
+
+    with op.batch_alter_table('user', schema=None) as batch_op:
+        batch_op.drop_column('pw')
+        batch_op.alter_column('a1',
+               existing_type=sa.TEXT(),
+               server_default='x',
+               existing_nullable=True)
+        batch_op.alter_column('name',
+               existing_type=sa.VARCHAR(length=50),
+               nullable=False)
+
+    ### end Alembic commands ###""")
+
+        eq_(re.sub(r"u'", "'", template_args['downgrades']),
+            """### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('user', schema=None) as batch_op:
+        batch_op.alter_column('name',
+               existing_type=sa.VARCHAR(length=50),
+               nullable=True)
+        batch_op.alter_column('a1',
+               existing_type=sa.TEXT(),
+               server_default=None,
+               existing_nullable=True)
+        batch_op.add_column(sa.Column('pw', sa.VARCHAR(length=50), nullable=True))
+
+    with op.batch_alter_table('order', schema=None) as batch_op:
+        batch_op.alter_column('amount',
+               existing_type=sa.Numeric(precision=10, scale=2),
+               type_=sa.NUMERIC(precision=8, scale=2),
+               nullable=False,
+               existing_server_default=sa.text('0'))
+        batch_op.drop_column('user_id')
+
+    with op.batch_alter_table('address', schema=None) as batch_op:
+        batch_op.drop_column('street')
+
+    op.create_table('extra',
+    sa.Column('x', sa.CHAR(), nullable=True),
+    sa.Column('uid', sa.INTEGER(), nullable=True),
+    sa.ForeignKeyConstraint(['uid'], ['user.id'], )
+    )
+    op.drop_table('item')
+    ### end Alembic commands ###""")
+
     def test_include_symbol(self):
         context = MigrationContext.configure(
             connection=self.bind.connect(),