1 files changed, 37 insertions, 8 deletions

diff --git a/src/mailman_pgp/testing/layers.py b/src/mailman_pgp/testing/layers.py
index aabf216..fb8a3ec 100644
--- a/src/mailman_pgp/testing/layers.py
+++ b/src/mailman_pgp/testing/layers.py
@@ -14,6 +14,7 @@
 #
 # You should have received a copy of the GNU General Public License along with
 # this program.  If not, see <http://www.gnu.org/licenses/>.
+import contextlib
 import os
 from os.path import isfile
 
@@ -24,18 +25,36 @@ from mailman_pgp.database import transaction
 from mailman_pgp.model.base import Base
 
 
-def reset_pgp_world():
+def reset_rollback():
+    config.db.session.rollback()
+
+
+def reset_pgp_dirs():
     for keydir in (config.pgp.keydir_config.values()):
         for path in os.listdir(keydir):
             full_path = os.path.join(keydir, path)
             if isfile(full_path):
                 os.remove(full_path)
-    config.db.session.rollback()
+
+
+def reset_pgp_hard():
+    reset_rollback()
+    reset_pgp_dirs()
     with transaction():
         Base.metadata.drop_all(config.db.engine)
         Base.metadata.create_all(config.db.engine)
 
 
+def reset_pgp_soft():
+    reset_rollback()
+    reset_pgp_dirs()
+    with contextlib.closing(config.db.engine.connect()) as con:
+        trans = con.begin()
+        for table in reversed(Base.metadata.sorted_tables):
+            con.execute(table.delete())
+        trans.commit()
+
+
 # It's weird that ws have to do this, but for some reason nose2 test layers
 # don't work when ws create a mixin class with the two classmethods
 # and subclass both it and the respective Mailman Core test layer.
@@ -46,28 +65,38 @@ class PGPConfigLayer(ConfigLayer):
 
     @classmethod
     def tearDown(cls):
-        reset_pgp_world()
+        reset_pgp_soft()
+
+    @classmethod
+    def testTearDown(cls):
+        reset_pgp_soft()
+
+
+class PGPMigrationLayer(ConfigLayer):
+    @classmethod
+    def tearDown(cls):
+        reset_pgp_hard()
 
     @classmethod
     def testTearDown(cls):
-        reset_pgp_world()
+        reset_pgp_hard()
 
 
 class PGPSMTPLayer(SMTPLayer):
     @classmethod
     def tearDown(cls):
-        reset_pgp_world()
+        reset_pgp_soft()
 
     @classmethod
     def testTearDown(cls):
-        reset_pgp_world()
+        reset_pgp_soft()
 
 
 class PGPRESTLayer(RESTLayer):
     @classmethod
     def tearDown(cls):
-        reset_pgp_world()
+        reset_pgp_soft()
 
     @classmethod
     def testTearDown(cls):
-        reset_pgp_world()
+        reset_pgp_soft()