diff --git a/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py b/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py index 6d467c2995636..e7c2496718d75 100644 --- a/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py +++ b/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py @@ -106,37 +106,37 @@ def hook(self) -> SmtpHook: def notify(self, context): """Send a email via smtp server.""" - fields_to_re_render = [] - if self.from_email is None: - if self.hook.from_email is not None: - self.from_email = self.hook.from_email - else: - raise ValueError("You should provide `from_email` or define it in the connection") - fields_to_re_render.append("from_email") - if self.subject is None: - smtp_default_templated_subject_path: str - if self.hook.subject_template: - smtp_default_templated_subject_path = self.hook.subject_template - else: - smtp_default_templated_subject_path = ( - Path(__file__).parent / "templates" / "email_subject.jinja2" - ).as_posix() - self.subject = self._read_template(smtp_default_templated_subject_path) - fields_to_re_render.append("subject") - if self.html_content is None: - smtp_default_templated_html_content_path: str - if self.hook.html_content_template: - smtp_default_templated_html_content_path = self.hook.html_content_template - else: - smtp_default_templated_html_content_path = ( - Path(__file__).parent / "templates" / "email.html" - ).as_posix() - self.html_content = self._read_template(smtp_default_templated_html_content_path) - fields_to_re_render.append("html_content") - if fields_to_re_render: - jinja_env = self.get_template_env(dag=context["dag"]) - self._do_render_template_fields(self, fields_to_re_render, context, jinja_env, set()) with self.hook as smtp: + fields_to_re_render = [] + if self.from_email is None: + if smtp.from_email is not None: + self.from_email = smtp.from_email + else: + raise ValueError("You should provide `from_email` or define it in the connection") + fields_to_re_render.append("from_email") + if self.subject is None: + smtp_default_templated_subject_path: str + if smtp.subject_template: + smtp_default_templated_subject_path = smtp.subject_template + else: + smtp_default_templated_subject_path = ( + Path(__file__).parent / "templates" / "email_subject.jinja2" + ).as_posix() + self.subject = self._read_template(smtp_default_templated_subject_path) + fields_to_re_render.append("subject") + if self.html_content is None: + smtp_default_templated_html_content_path: str + if smtp.html_content_template: + smtp_default_templated_html_content_path = smtp.html_content_template + else: + smtp_default_templated_html_content_path = ( + Path(__file__).parent / "templates" / "email.html" + ).as_posix() + self.html_content = self._read_template(smtp_default_templated_html_content_path) + fields_to_re_render.append("html_content") + if fields_to_re_render: + jinja_env = self.get_template_env(dag=context["dag"]) + self._do_render_template_fields(self, fields_to_re_render, context, jinja_env, set()) smtp.send_email_smtp( smtp_conn_id=self.smtp_conn_id, from_email=self.from_email, diff --git a/providers/smtp/tests/unit/smtp/notifications/test_smtp.py b/providers/smtp/tests/unit/smtp/notifications/test_smtp.py index ff3c02c4da956..73c1de253a5f5 100644 --- a/providers/smtp/tests/unit/smtp/notifications/test_smtp.py +++ b/providers/smtp/tests/unit/smtp/notifications/test_smtp.py @@ -124,8 +124,8 @@ def test_notifier_with_defaults(self, mock_smtphook_hook, create_task_instance): from_email="any email", to="test_reciver@test.com", ) - mock_smtphook_hook.return_value.subject_template = None - mock_smtphook_hook.return_value.html_content_template = None + mock_smtphook_hook.return_value.__enter__.return_value.subject_template = None + mock_smtphook_hook.return_value.__enter__.return_value.html_content_template = None notifier(context) mock_smtphook_hook.return_value.__enter__().send_email_smtp.assert_called_once_with( from_email="any email", @@ -157,10 +157,9 @@ def test_notifier_with_nondefault_connection_extra(self, mock_smtphook_hook, cre f_content.write("Mock content goes here") f_content.flush() - - mock_smtphook_hook.return_value.from_email = "{{ ti.task_id }}@test.com" - mock_smtphook_hook.return_value.subject_template = f_subject.name - mock_smtphook_hook.return_value.html_content_template = f_content.name + mock_smtphook_hook.return_value.__enter__.return_value.from_email = "{{ ti.task_id }}@test.com" + mock_smtphook_hook.return_value.__enter__.return_value.subject_template = f_subject.name + mock_smtphook_hook.return_value.__enter__.return_value.html_content_template = f_content.name notifier = SmtpNotifier( to="test_reciver@test.com", )