Merge pull request #33 from uber/mk_attachments_in_emails

Add starttls and attachment support to mail.

clay/mail.py

@@ -1,5 +1,7 @@
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
+from email.mime.base import MIMEBase
+from email import Encoders
 import smtplib
 import six
 
@@ -21,7 +23,8 @@ def _string_or_list(obj):
         return obj
 
 
-def sendmail(mailto, subject, message, subtype='html', charset='utf-8', smtpconfig=None, **headers):
+def sendmail(mailto, subject, message, subtype='html', charset='utf-8',
+             smtpconfig=None, attachments={}, use_starttls=False, **headers):
     '''
     Send an email to the given address. Additional SMTP headers may be specified
     as keyword arguments.
@@ -45,6 +48,17 @@ def sendmail(mailto, subject, message, subtype='html', charset='utf-8', smtpconf
     text = MIMEText(message, subtype, charset)
     msg.attach(text)
 
+    # Add attachments
+    for file_name, file_payload in attachments.items():
+        part = MIMEBase('application', 'octet-stream')
+        part.set_payload(file_payload.encode(charset))
+        Encoders.encode_base64(part)
+        part.add_header(
+            'Content-Disposition',
+            'attachment; filename="%s"' % file_name
+        )
+        msg.attach(part)
+
     if not 'From' in msg:
         msg['From'] = smtpconfig.get('from')
     mailfrom = msg['From']
@@ -58,6 +72,10 @@ def sendmail(mailto, subject, message, subtype='html', charset='utf-8', smtpconf
 
     smtp = smtplib.SMTP(smtpconfig.get('host'), smtpconfig.get('port'))
     if smtpconfig.get('username', None) is not None and smtpconfig.get('password', None) is not None:
+        if use_starttls:
+            smtp.elho()
+            smtp.starttls()
+            smtp.elho()
         smtp.login(smtpconfig.get('username'), smtpconfig.get('password'))
     smtp.sendmail(mailfrom, recipients, msg.as_string())
     smtp.quit()

tests/test_mail.py

@@ -29,6 +29,8 @@ def test_sendmail(self, mock_SMTP):
         self.assertIn('From: %s' % from_header, args[2])
         self.assertIn('Subject: %s' % subject, args[2])
         self.assertIn('Content-Type: text/html', args[2])
+        self.assertEqual(len(mock_SMTP_instance.elho.call_args_list), 0)
+        self.assertFalse(mock_SMTP_instance.starttls.called)
 
     @mock.patch("smtplib.SMTP")
     def test_sendmail_with_other_smtpconfig(self, mock_SMTP):
@@ -51,3 +53,36 @@ def test_sendmail_with_other_smtpconfig(self, mock_SMTP):
         self.assertIn('From: %s' % from_header, args[2])
         self.assertIn('Subject: %s' % subject, args[2])
         self.assertIn('Content-Type: text/html', args[2])
+
+    @mock.patch('smtplib.SMTP')
+    def test_sendmail_starttls(self, mock_SMTP):
+        """Test that when use_starttls is true, elho and starttls are called."""
+        mock_SMTP_instance = mock_SMTP.return_value
+
+        mailto = 'otherfake@email.com'
+        subject = 'This is another subject'
+        message = 'This is another message'
+        mail.sendmail(mailto, subject, message, use_starttls=True)
+
+        self.assertEqual(len(mock_SMTP_instance.elho.call_args_list), 2)
+        self.assertTrue(mock_SMTP_instance.starttls.called)
+
+    @mock.patch('smtplib.SMTP')
+    @mock.patch('email.mime.base.MIMEBase.set_payload')
+    def test_sendmail_attachments(self, mock_set_payload, mock_SMTP):
+        """Test that whens sarttls is true, elho and starttls are called."""
+
+        mailto = 'otherfake@email.com'
+        subject = 'This is another subject'
+        message = 'This is another message'
+        file_name = 'my_file_name.txt'
+        file_content = 'this is the content'
+        mail.sendmail(
+            mailto,
+            subject,
+            message,
+            attachments={file_name: file_content}
+        )
+
+        # set_payload should be called on the file content.
+        self.assertIn(file_content, [call[0][0] for call in mock_set_payload.call_args_list])