Skip to content

Commit

Permalink
Update tests and provider version pin
Browse files Browse the repository at this point in the history
  • Loading branch information
vchiapaikeo committed Feb 26, 2024
1 parent db6e0f6 commit cb466f6
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 28 deletions.
31 changes: 15 additions & 16 deletions airflow/providers/smtp/notifications/smtp.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,23 +25,12 @@
from airflow.notifications.basenotifier import BaseNotifier
from airflow.providers.smtp.hooks.smtp import SmtpHook

SMTP_DEFAULT_TEMPLATED_HTML_CONTENT_PATH = conf.get(
"smtp",
"templated_html_content_path",
fallback=(Path(__file__).parent / "templates" / "email.html").as_posix(),
)
SMTP_DEFAULT_TEMPLATED_SUBJECT_PATH = conf.get(
"smtp",
"templated_email_subject_path",
fallback=(Path(__file__).parent / "templates" / "email_subject.jinja2").as_posix(),
)


class SmtpNotifier(BaseNotifier):
"""
SMTP Notifier.
Accepts keyword arguments. The only required argument is `to`. Examples:
Accepts keyword arguments. The only required arguments are `from_email` and `to`. Examples:
.. code-block:: python
Expand Down Expand Up @@ -100,24 +89,34 @@ def __init__(
self.smtp_conn_id = smtp_conn_id
self.from_email = from_email or conf.get("smtp", "smtp_mail_from")
self.to = to
self.subject = (
subject or Path(SMTP_DEFAULT_TEMPLATED_SUBJECT_PATH).read_text().replace("\n", "").strip()
)
self.files = files
self.cc = cc
self.bcc = bcc
self.mime_subtype = mime_subtype
self.mime_charset = mime_charset
self.custom_headers = custom_headers

smtp_default_templated_subject_path = conf.get(
"smtp",
"templated_email_subject_path",
fallback=(Path(__file__).parent / "templates" / "email_subject.jinja2").as_posix(),
)
self.subject = (
subject or Path(smtp_default_templated_subject_path).read_text().replace("\n", "").strip()
)
# If html_content is passed, prioritize it. Otherwise, if template is passed, use
# it to populate html_content. Else, fall back to defaults defined in settings
if html_content is not None:
self.html_content = html_content
elif template is not None:
self.html_content = Path(template).read_text()
else:
self.html_content = Path(SMTP_DEFAULT_TEMPLATED_HTML_CONTENT_PATH).read_text()
smtp_default_templated_html_content_path = conf.get(
"smtp",
"templated_html_content_path",
fallback=(Path(__file__).parent / "templates" / "email.html").as_posix(),
)
self.html_content = Path(smtp_default_templated_html_content_path).read_text()

@cached_property
def hook(self) -> SmtpHook:
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/smtp/provider.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -92,14 +92,14 @@ config:
Allows overriding of the standard templated email subject line when the SmtpNotifier is used.
Must provide a path.
type: string
version_added: 2.9.0
version_added: 1.6.1
example: "path/to/override/email_subject.html"
default: ~
default_templated_html_content_path:
description: |
Allows overriding of the standard templated email path when the SmtpNotifier is used. Must provide
a path.
type: string
version_added: 2.9.0
version_added: 1.6.1
example: "path/to/override/email.html"
default: ~
11 changes: 1 addition & 10 deletions airflow/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@

from airflow import policies
from airflow.configuration import AIRFLOW_HOME, WEBSERVER_CONFIG, conf # NOQA F401
from airflow.exceptions import AirflowProviderDeprecationWarning, RemovedInAirflow3Warning
from airflow.exceptions import RemovedInAirflow3Warning
from airflow.executors import executor_constants
from airflow.logging_config import configure_logging
from airflow.utils.orm_event_handlers import setup_event_handlers
Expand Down Expand Up @@ -503,15 +503,6 @@ def import_local_settings():
setattr(airflow_local_settings, "task_policy", airflow_local_settings.policy)
names.remove("policy")

if "SMTP_DEFAULT_TEMPLATED_SUBJECT" in names or "SMTP_DEFAULT_TEMPLATED_HTML_CONTENT_PATH" in names:
warnings.warn(
"Configuring non-default `SMTP_DEFAULT_TEMPLATED_SUBJECT` and "
"`SMTP_DEFAULT_TEMPLATED_HTML_CONTENT_PATH` is deprecated. Please upgrade to "
"the new version of the SMTP provider and use provider configurations instead.",
AirflowProviderDeprecationWarning,
stacklevel=2,
)

plugin_functions = policies.make_plugin_from_local_settings(
POLICY_PLUGIN_MANAGER, airflow_local_settings, names
)
Expand Down
41 changes: 41 additions & 0 deletions tests/providers/smtp/notifications/test_smtp.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

from __future__ import annotations

import tempfile
from unittest import mock

import pytest
Expand All @@ -30,6 +31,7 @@
send_smtp_notification,
)
from airflow.utils import timezone
from tests.test_utils.config import conf_vars

pytestmark = pytest.mark.db_test

Expand Down Expand Up @@ -166,3 +168,42 @@ def test_notifier_with_defaults_sla(self, mock_smtphook_hook, dag_maker):
mime_charset="utf-8",
custom_headers=None,
)

@mock.patch("airflow.providers.smtp.notifications.smtp.SmtpHook")
def test_notifier_with_nondefault_conf_vars(self, mock_smtphook_hook, create_task_instance):
ti = create_task_instance(dag_id="dag", task_id="op", execution_date=timezone.datetime(2018, 1, 1))
context = {"dag": ti.dag_run.dag, "ti": ti}

with tempfile.NamedTemporaryFile(mode="wt", suffix=".txt") as f_subject, tempfile.NamedTemporaryFile(
mode="wt", suffix=".txt"
) as f_content:
f_subject.write("Task {{ ti.task_id }} failed")
f_subject.flush()

f_content.write("Mock content goes here")
f_content.flush()

with conf_vars(
{
("smtp", "templated_html_content_path"): f_content.name,
("smtp", "templated_email_subject_path"): f_subject.name,
}
):
notifier = SmtpNotifier(
from_email=conf.get("smtp", "smtp_mail_from"),
to="[email protected]",
)
notifier(context)
mock_smtphook_hook.return_value.__enter__().send_email_smtp.assert_called_once_with(
from_email=conf.get("smtp", "smtp_mail_from"),
to="[email protected]",
subject="Task op failed",
html_content="Mock content goes here",
smtp_conn_id="smtp_default",
files=None,
cc=None,
bcc=None,
mime_subtype="mixed",
mime_charset="utf-8",
custom_headers=None,
)

0 comments on commit cb466f6

Please sign in to comment.