def test_execute(self, mock_imap_hook, mock_s3_hook):
mock_imap_hook.return_value.__enter__ = mock_imap_hook
mock_imap_hook.return_value.retrieve_mail_attachments.return_value = [('test_file', b'Hello World')]
ImapAttachmentToS3Operator(**self.kwargs).execute(context={})
mock_imap_hook.return_value.retrieve_mail_attachments.assert_called_once_with(
name=self.kwargs['imap_attachment_name'],
check_regex=self.kwargs['imap_check_regex'],
latest_only=True,
mail_folder=self.kwargs['imap_mail_folder'],
mail_filter=self.kwargs['imap_mail_filter'],
)
mock_s3_hook.return_value.load_bytes.assert_called_once_with(
bytes_data=mock_imap_hook.return_value.retrieve_mail_attachments.return_value[0][1],
key=self.kwargs['s3_key'],
replace=self.kwargs['s3_overwrite'],
)
from datetime import datetime
from os import getenv
from airflow import DAG
from airflow.providers.amazon.aws.transfers.imap_attachment_to_s3 import ImapAttachmentToS3Operator
# [START howto_operator_imap_attachment_to_s3_env_variables]
IMAP_ATTACHMENT_NAME = getenv("IMAP_ATTACHMENT_NAME", "test.txt")
IMAP_MAIL_FOLDER = getenv("IMAP_MAIL_FOLDER", "INBOX")
IMAP_MAIL_FILTER = getenv("IMAP_MAIL_FILTER", "All")
S3_DESTINATION_KEY = getenv("S3_DESTINATION_KEY", "s3://bucket/key.json")
# [END howto_operator_imap_attachment_to_s3_env_variables]
with DAG(
dag_id="example_imap_attachment_to_s3",
start_date=datetime(2021, 1, 1),
schedule_interval=None,
catchup=False,
tags=['example'],
) as dag:
# [START howto_operator_imap_attachment_to_s3_task_1]
task_transfer_imap_attachment_to_s3 = ImapAttachmentToS3Operator(
imap_attachment_name=IMAP_ATTACHMENT_NAME,
s3_key=S3_DESTINATION_KEY,
imap_mail_folder=IMAP_MAIL_FOLDER,
imap_mail_filter=IMAP_MAIL_FILTER,
task_id='transfer_imap_attachment_to_s3',
)
# [END howto_operator_imap_attachment_to_s3_task_1]
from datetime import datetime
from os import getenv
from airflow import DAG
from airflow.providers.amazon.aws.transfers.imap_attachment_to_s3 import ImapAttachmentToS3Operator
IMAP_ATTACHMENT_NAME = getenv("IMAP_ATTACHMENT_NAME", "test.txt")
IMAP_MAIL_FOLDER = getenv("IMAP_MAIL_FOLDER", "INBOX")
IMAP_MAIL_FILTER = getenv("IMAP_MAIL_FILTER", "All")
S3_BUCKET = getenv("S3_BUCKET", "test-bucket")
S3_KEY = getenv("S3_KEY", "key")
with DAG(
dag_id="example_imap_attachment_to_s3",
start_date=datetime(2021, 1, 1),
schedule_interval=None,
catchup=False,
tags=['example'],
) as dag:
# [START howto_transfer_imap_attachment_to_s3]
task_transfer_imap_attachment_to_s3 = ImapAttachmentToS3Operator(
task_id='transfer_imap_attachment_to_s3',
imap_attachment_name=IMAP_ATTACHMENT_NAME,
s3_bucket=S3_BUCKET,
s3_key=S3_KEY,
imap_mail_folder=IMAP_MAIL_FOLDER,
imap_mail_filter=IMAP_MAIL_FILTER,
)
# [END howto_transfer_imap_attachment_to_s3]