Fix StepCollector caching causing multiple copies of steps to be registered

2 files changed, 27 insertions, 10 deletions

diff --git a/yoyo/migrations.py b/yoyo/migrations.py
index 5cd9741..9e84900 100755
--- a/yoyo/migrations.py
+++ b/yoyo/migrations.py
@@ -27,7 +27,6 @@ from yoyo.utils import plural
 
 logger = getLogger('yoyo.migrations')
 default_migration_table = '_yoyo_migration'
-_step_collectors = defaultdict(lambda: StepCollector())
 
 
 class Migration(object):
@@ -63,10 +62,11 @@ class Migration(object):
             self.source = source = f.read()
             migration_code = compile(source, f.name, 'exec')
 
-        collector = _step_collectors[f.name]
+        collector = StepCollector()
         ns = {'step': collector.add_step,
               'group': collector.add_step_group,
-              'transaction': collector.add_step_group}
+              'transaction': collector.add_step_group,
+              'collector': collector}
         try:
             exec_(migration_code, ns)
         except Exception as e:
@@ -293,7 +293,6 @@ def read_migrations(*directories):
                 migrations.post_apply.append(migration)
             else:
                 migrations.append(migration)
-
     return migrations
 
 
@@ -412,18 +411,16 @@ class StepCollector(object):
 
 
 def _get_collector(depth=2):
-    fi = inspect.getframeinfo(inspect.stack()[depth][0])
-    return _step_collectors[fi.filename]
+    return inspect.stack()[depth][0].f_locals['collector']
 
 
 def step(*args, **kwargs):
-    collector = _get_collector()
-    return collector.add_step(*args, **kwargs)
+    return _get_collector().add_step(*args, **kwargs)
 
 
 def group(*args, **kwargs):
-    collector = _get_collector()
-    return collector.add_step_group(*args, **kwargs)
+    return _get_collector().add_step_group(*args, **kwargs)
+
 
 #: Alias for compatibility purposes.
 #: This no longer affects transaction handling.
diff --git a/yoyo/tests/test_migrations.py b/yoyo/tests/test_migrations.py
index dad6aae..6ac8bba 100644
--- a/yoyo/tests/test_migrations.py
+++ b/yoyo/tests/test_migrations.py
@@ -347,3 +347,23 @@ class TestReadMigrations(object):
         migrations = read_migrations(tmpdir)
         assert len(migrations) == 0
         assert len(migrations.post_apply) == 1
+
+    @with_migrations(**{'a': '''step('SELECT 1')'''})
+    def test_it_does_not_add_duplicate_steps(self, tmpdir):
+        m = read_migrations(tmpdir)[0]
+        m.load()
+        assert len(m.steps) == 1
+
+        m = read_migrations(tmpdir)[0]
+        m.load()
+        assert len(m.steps) == 1
+
+    @with_migrations(**{'a': '''from yoyo import step; step('SELECT 1')'''})
+    def test_it_does_not_add_duplicate_steps_with_imported_symbols(self, tmpdir):
+        m = read_migrations(tmpdir)[0]
+        m.load()
+        assert len(m.steps) == 1
+
+        m = read_migrations(tmpdir)[0]
+        m.load()
+        assert len(m.steps) == 1