From 91d4c00a84d1a75dd7ede4eadcd2e6ff3431e11c Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Sat, 4 Jan 2014 13:25:25 +0200
Subject: [PATCH] Fix generates decorator

---
 sqlalchemy_utils/decorators.py | 26 +++++++++++++++-----------
 tests/test_generates.py        |  4 ----
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/sqlalchemy_utils/decorators.py b/sqlalchemy_utils/decorators.py
index 7fbc293..e10c2fc 100644
--- a/sqlalchemy_utils/decorators.py
+++ b/sqlalchemy_utils/decorators.py
@@ -6,11 +6,23 @@ import six
 
 class AttributeValueGenerator(object):
     def __init__(self):
+        self.listener_args = [
+            (
+                sa.orm.mapper,
+                'mapper_configured',
+                self.update_generator_registry
+            ),
+            (
+                sa.orm.session.Session,
+                'before_flush',
+                self.update_generated_properties
+            )
+        ]
         self.reset()
 
     def reset(self):
-        self.generator_registry = defaultdict(list)
         self.listeners_registered = False
+        self.generator_registry = defaultdict(list)
 
     def generator_wrapper(self, func, attr):
         def wrapper(self, *args, **kwargs):
@@ -29,16 +41,8 @@ class AttributeValueGenerator(object):
 
     def register_listeners(self):
         if not self.listeners_registered:
-            sa.event.listen(
-                sa.orm.mapper,
-                'mapper_configured',
-                self.update_generator_registry
-            )
-            sa.event.listen(
-                sa.orm.session.Session,
-                'before_flush',
-                self.update_generated_properties
-            )
+            for args in self.listener_args:
+                sa.event.listen(*args)
             self.listeners_registered = True
 
     def update_generator_registry(self, mapper, class_):
diff --git a/tests/test_generates.py b/tests/test_generates.py
index 29aaedf..a8a7c38 100644
--- a/tests/test_generates.py
+++ b/tests/test_generates.py
@@ -5,10 +5,6 @@ from tests import TestCase
 
 
 class GeneratesTestCase(TestCase):
-    def teardown_method(self, method):
-        TestCase.teardown_method(self, method)
-        generator.reset()
-
     def test_generates_value_before_flush(self):
         article = self.Article()
         article.name = u'some article name'