2 files changed, 19 insertions, 8 deletions

diff --git a/src/mailman_pgp/workflows/base.py b/src/mailman_pgp/workflows/base.py
index 8a8126a..4ca1525 100644
--- a/src/mailman_pgp/workflows/base.py
+++ b/src/mailman_pgp/workflows/base.py
@@ -18,6 +18,7 @@
 """"""
 from mailman.email.message import UserNotification
 from mailman.interfaces.subscriptions import TokenOwner
+from mailman.workflows.common import WhichSubscriber
 from pgpy import PGPKey
 
 from mailman_pgp.model.address import PGPAddress
@@ -63,11 +64,12 @@ class PubkeyMixin:
             self.pubkey = None
 
     def _step_pubkey_checks(self):
+        self.push('restore_subscriber')
         if not self.pubkey:
             self.push('send_key_request')
         else:
             if not self.pubkey_confirmed:
-                self.push('send_confirm_request')
+                self.push('send_key_confirm_request')
 
     def _step_send_key_request(self):
         self._set_token(TokenOwner.subscriber)
@@ -92,11 +94,11 @@ class PubkeyMixin:
         else:
             self.pubkey = pgp_address.key
             if not self.pubkey_confirmed:
-                self.push('send_confirm_request')
+                self.push('send_key_confirm_request')
 
-    def _step_send_confirm_request(self):
+    def _step_send_key_confirm_request(self):
         self._set_token(TokenOwner.subscriber)
-        self.push('receive_confirmation')
+        self.push('receive_key_confirmation')
         self.save()
         request_address = self.mlist.request_address
         email_address = self.address.email
@@ -114,5 +116,12 @@ class PubkeyMixin:
         msg.send(self.mlist)
         raise StopIteration
 
-    def _step_receive_confirmation(self):
+    def _step_receive_key_confirmation(self):
         self._set_token(TokenOwner.no_one)
+
+    def _step_restore_subscriber(self):
+        if self.which is WhichSubscriber.address:
+            self.subscriber = self.address
+        else:
+            assert self.which is WhichSubscriber.user
+            self.subscriber = self.user
diff --git a/src/mailman_pgp/workflows/tests/test_base.py b/src/mailman_pgp/workflows/tests/test_base.py
index a52571c..c99c6d6 100644
--- a/src/mailman_pgp/workflows/tests/test_base.py
+++ b/src/mailman_pgp/workflows/tests/test_base.py
@@ -88,6 +88,7 @@ class TestPubkeyMixin(unittest.TestCase):
                                              pubkey=self.sender_key.pubkey,
                                              pubkey_pre_confirmed=True)
         workflow.run_thru('pubkey_checks')
+        next(workflow)
         with patch.object(workflow, '_step_do_subscription') as step:
             next(workflow)
         step.assert_called_once_with()
@@ -112,7 +113,7 @@ class TestPubkeyMixin(unittest.TestCase):
                          self.sender_key.pubkey.fingerprint)
 
         with patch.object(receive_workflow,
-                          '_step_send_confirm_request') as step:
+                          '_step_send_key_confirm_request') as step:
             next(receive_workflow)
         step.assert_called_once_with()
 
@@ -147,6 +148,7 @@ class TestPubkeyMixin(unittest.TestCase):
         receive_workflow.token = workflow.token
         receive_workflow.restore()
         receive_workflow.run_thru('receive_key')
+        next(receive_workflow)
         with patch.object(receive_workflow, '_step_do_subscription') as step:
             next(receive_workflow)
         step.assert_called_once_with()
@@ -182,8 +184,8 @@ class TestPubkeyMixin(unittest.TestCase):
         receive_workflow = ConfirmSubscriptionPolicy(self.mlist)
         receive_workflow.token = workflow.token
         receive_workflow.restore()
-        receive_workflow.run_thru('receive_confirmation')
-
+        receive_workflow.run_thru('receive_key_confirmation')
+        next(receive_workflow)
         with patch.object(receive_workflow, '_step_do_subscription') as step:
             next(receive_workflow)
         step.assert_called_once_with()