Added decorator to clean up the adaptation mappings after testsfix-948

Many tests were doing it manually, some weren't doing it and resulted in
failure if run in different order.

Close #948

1 files changed, 17 insertions, 0 deletions

diff --git a/tests/testutils.py b/tests/testutils.py
index 405da18..996ceb7 100644
--- a/tests/testutils.py
+++ b/tests/testutils.py
@@ -429,3 +429,20 @@ def slow(f):
             return self.skipTest("slow test")
         return f(self)
     return slow_
+
+
+def restore_types(f):
+    """Decorator to restore the adaptation system after running a test"""
+    @wraps(f)
+    def restore_types_(self):
+        types = psycopg2.extensions.string_types.copy()
+        adapters = psycopg2.extensions.adapters.copy()
+        try:
+            return f(self)
+        finally:
+            psycopg2.extensions.string_types.clear()
+            psycopg2.extensions.string_types.update(types)
+            psycopg2.extensions.adapters.clear()
+            psycopg2.extensions.adapters.update(adapters)
+
+    return restore_types_