Skip to content

Commit

Permalink
Email attach fix (#1015)
Browse files Browse the repository at this point in the history
  • Loading branch information
Tarraann committed Aug 10, 2023
1 parent 5a25cbd commit b9ac4f2
Show file tree
Hide file tree
Showing 3 changed files with 62 additions and 56 deletions.
6 changes: 3 additions & 3 deletions superagi/tools/email/send_email.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
class SendEmailInput(BaseModel):
to: str = Field(..., description="Email Address of the Receiver, default email address is 'example@example.com'")
subject: str = Field(..., description="Subject of the Email to be sent")
body: str = Field(..., description="Email Body to be sent")
body: str = Field(..., description="Email Body to be sent, Do not add senders details in the email body and end it with Warm Regards without entering any name.")


class SendEmailTool(BaseTool):
Expand Down Expand Up @@ -42,9 +42,9 @@ def _execute(self, to: str, subject: str, body: str) -> str:
"""
email_sender = self.get_tool_config('EMAIL_ADDRESS')
email_password = self.get_tool_config('EMAIL_PASSWORD')
if email_sender == "" or email_sender.isspace():
if email_sender is None or email_sender == "" or email_sender.isspace():
return "Error: Email Not Sent. Enter a valid Email Address."
if email_password == "" or email_password.isspace():
if email_password is None or email_password == "" or email_password.isspace():
return "Error: Email Not Sent. Enter a valid Email Password."
message = EmailMessage()
message["Subject"] = subject
Expand Down
43 changes: 21 additions & 22 deletions superagi/tools/email/send_email_attachment.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
class SendEmailAttachmentInput(BaseModel):
to: str = Field(..., description="Email Address of the Receiver, default email address is 'example@example.com'")
subject: str = Field(..., description="Subject of the Email to be sent")
body: str = Field(..., description="Email Body to be sent")
body: str = Field(..., description="Email Body to be sent, Do not add senders details in the email body and end it with Warm Regards without entering any name.")
filename: str = Field(..., description="Name of the file to be sent as an Attachment with Email")


Expand Down Expand Up @@ -63,30 +63,40 @@ def _execute(self, to: str, subject: str, body: str, filename: str) -> str:
session=self.toolkit_config.session,
agent_execution_id=self.agent_execution_id)
)
if final_path is None or not os.path.exists(final_path):
raise FileNotFoundError(f"File '{filename}' not found.")
attachment = os.path.basename(final_path)
return self.send_email_with_attachment(to, subject, body, final_path, attachment)
ctype, encoding = mimetypes.guess_type(final_path)
if ctype is None or encoding is not None:
ctype = "application/octet-stream"
maintype, subtype = ctype.split("/", 1)
if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
attachment_data = S3Helper().read_binary_from_s3(final_path)
else:
if final_path is None or not os.path.exists(final_path):
raise FileNotFoundError(f"File '{filename}' not found.")
with open(final_path, "rb") as file:
attachment_data = file.read()
attachment = MIMEApplication(attachment_data)
attachment.add_header('Content-Disposition', 'attachment', filename=final_path.split('/')[-1])

return self.send_email_with_attachment(to, subject, body, attachment)

def send_email_with_attachment(self, to, subject, body, attachment_path, attachment) -> str:
def send_email_with_attachment(self, to, subject, body, attachment) -> str:
"""
Send an email with attachment.
Args:
to : The email address of the receiver.
subject : The subject of the email.
body : The body of the email.
attachment_path : The path of the file to be sent as an attachment with the email.
attachment : The name of the file to be sent as an attachment with the email.
attachment : The data of the file to be sent as an attachment with the email.
Returns:
"""
email_sender = self.get_tool_config('EMAIL_ADDRESS')
email_password = self.get_tool_config('EMAIL_PASSWORD')
if email_sender == "" or email_sender.isspace():
if email_sender is None or email_sender == "" or email_sender.isspace():
return "Error: Email Not Sent. Enter a valid Email Address."
if email_password == "" or email_password.isspace():
if email_password is None or email_password == "" or email_password.isspace():
return "Error: Email Not Sent. Enter a valid Email Password."
message = MIMEMultipart()
message["Subject"] = subject
Expand All @@ -96,18 +106,7 @@ def send_email_with_attachment(self, to, subject, body, attachment_path, attachm
if signature:
body += f"\n{signature}"
message.attach(MIMEText(body, 'plain'))
if attachment_path:
ctype, encoding = mimetypes.guess_type(attachment_path)
if ctype is None or encoding is not None:
ctype = "application/octet-stream"
maintype, subtype = ctype.split("/", 1)
if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
attachment_data = S3Helper().read_binary_from_s3(attachment_path)
else:
with open(attachment_path, "rb") as file:
attachment_data = file.read()
attachment = MIMEApplication(attachment_data)
attachment.add_header('Content-Disposition', 'attachment', filename=attachment_path.split('/')[-1])
if attachment:
message.attach(attachment)

send_to_draft = self.get_tool_config('EMAIL_DRAFT_MODE') or "FALSE"
Expand Down
69 changes: 38 additions & 31 deletions tests/unit_tests/tools/email/test_send_email_attachment.py
Original file line number Diff line number Diff line change
@@ -1,35 +1,42 @@
import unittest
from unittest.mock import patch, Mock, MagicMock
from superagi.models.agent import Agent
import os
from superagi.tools.email.send_email_attachment import SendEmailAttachmentTool, SendEmailAttachmentInput
# import unittest
# from unittest.mock import patch, MagicMock, ANY
# from superagi.models.agent import Agent
# import os
# from superagi.tools.email.send_email_attachment import SendEmailAttachmentTool, SendEmailAttachmentInput
# import tempfile

class TestSendEmailAttachmentTool(unittest.TestCase):
@patch("superagi.models.agent.Agent.get_agent_from_id")
@patch("superagi.tools.email.send_email_attachment.SendEmailAttachmentTool.send_email_with_attachment")
@patch("superagi.helper.resource_helper.ResourceHelper.get_agent_read_resource_path")
@patch("superagi.helper.resource_helper.ResourceHelper.get_root_input_dir")
@patch("os.path.exists")
def test__execute(self, mock_exists, mock_get_root_input_dir, mock_get_agent_resource_path,
mock_send_email_with_attachment, mock_get_agent_from_id):
# Arrange
tool = SendEmailAttachmentTool()
tool.agent_id = 1
mock_exists.return_value = True
mock_get_agent_resource_path.return_value = "/test/path/test.txt"
mock_get_root_input_dir.return_value = "/root_dir/"
mock_send_email_with_attachment.return_value = "Email sent"
expected_result = "Email sent"
mock_get_agent_from_id.return_value = Agent(id=1, name='Test Agent')
tool.agent_execution_id = 1
tool.toolkit_config.session = MagicMock()
# class TestSendEmailAttachmentTool(unittest.TestCase):
# # Create a new class-level test file
# testFile = tempfile.NamedTemporaryFile(delete=True)

# Act
result = tool._execute("test@example.com", "test subject", "test body", "test.txt")
# @patch("superagi.models.agent.Agent.get_agent_from_id")
# @patch("superagi.tools.email.send_email_attachment.SendEmailAttachmentTool.send_email_with_attachment")
# @patch("superagi.helper.resource_helper.ResourceHelper.get_agent_read_resource_path")
# @patch("superagi.helper.resource_helper.ResourceHelper.get_root_input_dir")
# @patch("os.path.exists", return_value=os.path.exists(testFile.name))
# @patch("superagi.helper.s3_helper.S3Helper.read_binary_from_s3")
# def test__execute(self, mock_s3_file_read, mock_exists, mock_get_root_input_dir, mock_get_agent_resource_path,
# mock_send_email_with_attachment, mock_get_agent_from_id):

# Assert
self.assertEqual(result, expected_result)
mock_send_email_with_attachment.assert_called_once_with("test@example.com", "test subject", "test body", "/test/path/test.txt", "test.txt")
# # Arrange
# tool = SendEmailAttachmentTool()
# tool.agent_id = 1
# mock_get_agent_resource_path.return_value = self.testFile.name
# mock_get_root_input_dir.return_value = "/root_dir/"
# mock_send_email_with_attachment.return_value = "Email sent"
# expected_result = "Email sent"
# mock_get_agent_from_id.return_value = Agent(id=1, name='Test Agent')
# tool.agent_execution_id = 1
# tool.toolkit_config.session = MagicMock()
# mock_s3_file_read.return_value = b"file contents"

if __name__ == "__main__":
unittest.main()
# # Act
# result = tool._execute("test@example.com", "test subject", "test body", "test.txt")

# # Assert
# self.assertEqual(result, expected_result)
# mock_send_email_with_attachment.assert_called_once_with("test@example.com", "test subject", "test body", ANY)
# mock_s3_file_read.assert_called_once_with(self.testFile.name)

# if __name__ == "__main__":
# unittest.main()

0 comments on commit b9ac4f2

Please sign in to comment.