Fix SMTP async tests

1 files changed, 17 insertions, 15 deletions

diff --git a/src/leap/mail/smtp/gateway.py b/src/leap/mail/smtp/gateway.py
index 1a187cf..9d78474 100644
--- a/src/leap/mail/smtp/gateway.py
+++ b/src/leap/mail/smtp/gateway.py
@@ -204,22 +204,24 @@ class SMTPDelivery(object):
             signal(proto.SMTP_RECIPIENT_ACCEPTED_ENCRYPTED, user.dest.addrstr)
 
         def not_found(failure):
-            if failure.check(KeyNotFound):
-                # if key was not found, check config to see if will send anyway
-                if self._encrypted_only:
-                    signal(proto.SMTP_RECIPIENT_REJECTED, user.dest.addrstr)
-                    raise smtp.SMTPBadRcpt(user.dest.addrstr)
-                log.msg("Warning: will send an unencrypted message (because "
-                        "encrypted_only' is set to False).")
-                signal(
-                    proto.SMTP_RECIPIENT_ACCEPTED_UNENCRYPTED,
-                    user.dest.addrstr)
-            else:
-                return failure
-
-        d = self._km.get_key(address, OpenPGPKey)  # might raise KeyNotFound
+            failure.trap(KeyNotFound)
+
+            # if key was not found, check config to see if will send anyway
+            if self._encrypted_only:
+                signal(proto.SMTP_RECIPIENT_REJECTED, user.dest.addrstr)
+                raise smtp.SMTPBadRcpt(user.dest.addrstr)
+            log.msg("Warning: will send an unencrypted message (because "
+                    "encrypted_only' is set to False).")
+            signal(
+                proto.SMTP_RECIPIENT_ACCEPTED_UNENCRYPTED,
+                user.dest.addrstr)
+
+        def encrypt_func(_):
+            return lambda: EncryptedMessage(user, self._outgoing_mail)
+
+        d = self._km.get_key(address, OpenPGPKey)
         d.addCallbacks(found, not_found)
-        d.addCallback(lambda _: EncryptedMessage(user, self._outgoing_mail))
+        d.addCallback(encrypt_func)
         return d
 
     def validateFrom(self, helo, origin):