2 files changed, 31 insertions, 8 deletions
diff --git a/src/mailman/model/tests/test_workflow.py b/src/mailman/model/tests/test_workflow.py
index 6081e5b57..ccf618c2b 100644
--- a/src/mailman/model/tests/test_workflow.py
+++ b/src/mailman/model/tests/test_workflow.py
@@ -108,3 +108,23 @@ class TestWorkflow(unittest.TestCase):
         self._manager.save(name, token)
         workflow = self._manager.restore('ewe', 'fly')
         self.assertIsNone(workflow)
+
+    def test_restore_removes_record(self):
+        name = 'ant'
+        token = 'bee'
+        self.assertEqual(self._manager.count(), 0)
+        self._manager.save(name, token)
+        self.assertEqual(self._manager.count(), 1)
+        self._manager.restore(name, token)
+        self.assertEqual(self._manager.count(), 0)
+
+    def test_save_after_restore(self):
+        name = 'ant'
+        token = 'bee'
+        self.assertEqual(self._manager.count(), 0)
+        self._manager.save(name, token)
+        self.assertEqual(self._manager.count(), 1)
+        self._manager.restore(name, token)
+        self.assertEqual(self._manager.count(), 0)
+        self._manager.save(name, token)
+        self.assertEqual(self._manager.count(), 1)
diff --git a/src/mailman/model/workflow.py b/src/mailman/model/workflow.py
index d9f23f53b..6ac3fa76a 100644
--- a/src/mailman/model/workflow.py
+++ b/src/mailman/model/workflow.py
@@ -51,15 +51,18 @@ class WorkflowStateManager:
     @dbconnection
     def save(self, store, name, token, step=None, data=None):
         """See `IWorkflowStateManager`."""
-        state = store.query(WorkflowState).get((name, token))
-        if state is None:
-            state = WorkflowState(name=name, token=token, step=step, data=data)
-            store.add(state)
-        else:
-            state.step = step
-            state.data = data
+        state = WorkflowState(name=name, token=token, step=step, data=data)
+        store.add(state)
 
     @dbconnection
     def restore(self, store, name, token):
         """See `IWorkflowStateManager`."""
-        return store.query(WorkflowState).get((name, token))
+        state = store.query(WorkflowState).get((name, token))
+        if state is not None:
+            store.delete(state)
+        return state
+
+    @dbconnection
+    def count(self, store):
+        """See `IWorkflowStateManager`."""
+        return store.query(WorkflowState).count()