Replace union_update with update

Fixed issue where removed method ``union_update()`` was used when a
customized :class:`.MigrationScript` instance included entries in the
``.imports`` data member, raising an AttributeError.

Change-Id: Ia141db106bc3d57238e2fa6a546041fd573c0ea4
Fixes: #512

3 files changed, 23 insertions, 5 deletions

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index b54e6aa..15b5b6b 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -395,7 +395,7 @@ class RevisionContext(object):
             # renders
             autogen_context.imports = set()
             if migration_script.imports:
-                autogen_context.imports.union_update(migration_script.imports)
+                autogen_context.imports.update(migration_script.imports)
             render._render_python_into_templatevars(
                 autogen_context, migration_script, template_args
             )
diff --git a/docs/build/unreleased/512.rst b/docs/build/unreleased/512.rst
new file mode 100644
index 0000000..8201a53
--- /dev/null
+++ b/docs/build/unreleased/512.rst
@@ -0,0 +1,8 @@
+.. change::
+   :tags: bug, autogenerate
+   :tickets: 512
+
+   Fixed issue where removed method ``union_update()`` was used when a
+   customized :class:`.MigrationScript` instance included entries in the
+   ``.imports`` data member, raising an AttributeError.
+
diff --git a/tests/test_script_production.py b/tests/test_script_production.py
index 7ab46ab..af01a38 100644
--- a/tests/test_script_production.py
+++ b/tests/test_script_production.py
@@ -605,7 +605,7 @@ class ImportsTest(TestBase):
     def tearDown(self):
         clear_staging_env()
 
-    def _env_fixture(self, target_metadata):
+    def _env_fixture(self, target_metadata, **kw):
         self.engine = engine = _sqlite_file_db()
 
         def run_env(self):
@@ -614,7 +614,8 @@ class ImportsTest(TestBase):
             with engine.connect() as connection:
                 context.configure(
                     connection=connection,
-                    target_metadata=target_metadata)
+                    target_metadata=target_metadata,
+                    **kw)
                 with context.begin_transaction():
                     context.run_migrations()
 
@@ -636,13 +637,22 @@ class ImportsTest(TestBase):
             Column('x', type_)
         )
 
-        with self._env_fixture(m):
+        def process_revision_directives(context, rev, generate_revisions):
+            generate_revisions[0].imports.add(
+                "from sqlalchemy.dialects.mysql import TINYINT")
+
+        with self._env_fixture(
+                m,
+                process_revision_directives=process_revision_directives
+        ):
             rev = command.revision(
                 self.cfg, message="some message",
                 autogenerate=True)
 
         with open(rev.path) as file_:
-            assert "from sqlalchemy.dialects import mysql" in file_.read()
+            contents = file_.read()
+            assert "from sqlalchemy.dialects import mysql" in contents
+            assert "from sqlalchemy.dialects.mysql import TINYINT" in contents
 
 
 class MultiContextTest(TestBase):