From da34963270ab4e2c0f9f2fde3a9733909ca5fa45 Mon Sep 17 00:00:00 2001 From: Shivam <6463385+shivaam@users.noreply.github.com> Date: Wed, 7 Oct 2026 07:59:33 -0700 Subject: [PATCH 1/2] Reconnect SMTP hooks after leaving a context manager Repeated notifications can reuse a cached hook whose SMTP session was closed by its previous context. Retaining that client makes the next context skip connection setup and attempt another send on a closed session. The hook must discard the client on context exit, including cleanup failures, so subsequent use can establish a fresh session. Keep cleanup exceptions visible rather than suppressing failures to make reuse work. Regression tests cover synchronous and asynchronous context reuse, close/quit failures, and repeated notification through the cached hook. --- .../src/airflow/providers/smtp/hooks/smtp.py | 10 ++++- .../smtp/tests/unit/smtp/hooks/test_smtp.py | 44 +++++++++++++++++++ .../unit/smtp/notifications/test_smtp.py | 26 +++++++++++ 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py b/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py index 744fe705cb70e..b8fbce819041d 100644 --- a/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py +++ b/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py @@ -82,11 +82,17 @@ async def __aenter__(self) -> SmtpHook: return await self.aget_conn() def __exit__(self, exc_type, exc_val, exc_tb): - self._smtp_client.close() + try: + self._smtp_client.close() + finally: + self._smtp_client = None async def __aexit__(self, exc_type, exc_val, exc_tb): if self._smtp_client: - await self._smtp_client.quit() + try: + await self._smtp_client.quit() + finally: + self._smtp_client = None def _setup_oauth2(self) -> tuple[str, str]: """ diff --git a/providers/smtp/tests/unit/smtp/hooks/test_smtp.py b/providers/smtp/tests/unit/smtp/hooks/test_smtp.py index e3a3f557f70ea..6d26ba94cec8b 100644 --- a/providers/smtp/tests/unit/smtp/hooks/test_smtp.py +++ b/providers/smtp/tests/unit/smtp/hooks/test_smtp.py @@ -22,6 +22,7 @@ import smtplib import ssl import tempfile +from contextlib import nullcontext from email.mime.application import MIMEApplication from unittest import mock from unittest.mock import AsyncMock, Mock, call, patch @@ -83,6 +84,27 @@ def _create_fake_smtp(mock_smtplib, use_ssl=True): class TestSmtpHook: + @pytest.mark.parametrize("close_fails", [False, True]) + @patch("smtplib.SMTP_SSL", autospec=True) + def test_reconnect_after_context_exit(self, mock_smtp_ssl, close_fails): + clients = [mock.create_autospec(smtplib.SMTP, instance=True) for _ in range(2)] + mock_smtp_ssl.side_effect = clients + if close_fails: + clients[0].close.side_effect = OSError("Close failed") + hook = SmtpHook() + + with pytest.raises(OSError, match="Close failed") if close_fails else nullcontext(): + with hook: + hook.send_email_smtp(to=TO_EMAIL, subject=TEST_SUBJECT, html_content=TEST_BODY) + + with hook: + hook.send_email_smtp(to=TO_EMAIL, subject=TEST_SUBJECT, html_content=TEST_BODY) + + assert mock_smtp_ssl.call_count == 2 + for client in clients: + client.sendmail.assert_called_once() + client.close.assert_called_once() + @pytest.fixture(autouse=True) def setup_connections(self, create_connection_without_db): create_connection_without_db( @@ -610,6 +632,28 @@ def test_test_connection_handles_noop_responses(self, mock_smtplib, noop_respons class TestSmtpHookAsync: """Tests for async functionality in SmtpHook.""" + @pytest.mark.parametrize("quit_fails", [False, True]) + async def test_reconnect_after_context_exit(self, mock_get_connection, mocker, quit_fails): + clients = [mock.create_autospec(aiosmtplib.SMTP, instance=True) for _ in range(2)] + mock_smtp = mocker.patch( + "airflow.providers.smtp.hooks.smtp.aiosmtplib.SMTP", autospec=True, side_effect=clients + ) + if quit_fails: + clients[0].quit.side_effect = OSError("Quit failed") + hook = SmtpHook() + + with pytest.raises(OSError, match="Quit failed") if quit_fails else nullcontext(): + async with hook: + await hook.asend_email_smtp(to=TO_EMAIL, subject=TEST_SUBJECT, html_content=TEST_BODY) + + async with hook: + await hook.asend_email_smtp(to=TO_EMAIL, subject=TEST_SUBJECT, html_content=TEST_BODY) + + assert mock_smtp.call_count == 2 + for client in clients: + client.sendmail.assert_awaited_once() + client.quit.assert_awaited_once() + @pytest.fixture(autouse=True) def setup_connections(self, create_connection_without_db): create_connection_without_db( diff --git a/providers/smtp/tests/unit/smtp/notifications/test_smtp.py b/providers/smtp/tests/unit/smtp/notifications/test_smtp.py index 4a5904405f146..305534bf00702 100644 --- a/providers/smtp/tests/unit/smtp/notifications/test_smtp.py +++ b/providers/smtp/tests/unit/smtp/notifications/test_smtp.py @@ -17,6 +17,8 @@ from __future__ import annotations +import json +import smtplib import tempfile from dataclasses import dataclass from unittest import mock @@ -24,6 +26,7 @@ import pytest +from airflow.providers.common.compat.sdk import Connection from airflow.providers.smtp.notifications.smtp import SmtpNotifier, send_smtp_notification from tests_common.test_utils.config import conf_vars @@ -94,6 +97,29 @@ def rendered(self) -> str: class TestSmtpNotifier: + @mock.patch("smtplib.SMTP_SSL", autospec=True) + def test_repeated_notifications_reconnect(self, mock_smtp_ssl, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=SMTP_CONN_ID, + conn_type="smtp", + host="smtp.example.com", + port=465, + extra=json.dumps({"disable_tls": True}), + ) + ) + clients = [mock.create_autospec(smtplib.SMTP, instance=True) for _ in range(2)] + mock_smtp_ssl.side_effect = clients + notifier = SmtpNotifier(**NOTIFIER_DEFAULT_PARAMS) + + notifier.notify({}) + notifier.notify({}) + + assert mock_smtp_ssl.call_count == 2 + for client in clients: + client.sendmail.assert_called_once() + client.close.assert_called_once() + @mock.patch("airflow.providers.smtp.notifications.smtp.SmtpHook") def test_notifier(_self, mock_smtphook_hook, create_dag_without_db): notifier = send_smtp_notification(**NOTIFIER_DEFAULT_PARAMS) From e1ea22ea45210762c91bb3ab00d808fbf2aa6c69 Mon Sep 17 00:00:00 2001 From: Shivam <6463385+shivaam@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:19:19 -0700 Subject: [PATCH 2/2] Keep SMTP reconnect tests compatible with older dependencies Use the existing async SMTP test double so lifecycle assertions do not enforce a newer third-party constructor signature. Set notification connection data as portable JSON instead of requiring SDK Connection.as_json on Airflow 3.0.6. Both failed CI environments are covered locally. All five corrected regressions still fail when context-exit client invalidation is removed, so the compatibility correction preserves their ability to catch the original bug. No production behavior or dependency floors are changed. --- .../smtp/tests/unit/smtp/hooks/test_smtp.py | 6 ++---- .../unit/smtp/notifications/test_smtp.py | 21 ++++++++++--------- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/providers/smtp/tests/unit/smtp/hooks/test_smtp.py b/providers/smtp/tests/unit/smtp/hooks/test_smtp.py index 6d26ba94cec8b..cab96aeaed9c4 100644 --- a/providers/smtp/tests/unit/smtp/hooks/test_smtp.py +++ b/providers/smtp/tests/unit/smtp/hooks/test_smtp.py @@ -634,10 +634,8 @@ class TestSmtpHookAsync: @pytest.mark.parametrize("quit_fails", [False, True]) async def test_reconnect_after_context_exit(self, mock_get_connection, mocker, quit_fails): - clients = [mock.create_autospec(aiosmtplib.SMTP, instance=True) for _ in range(2)] - mock_smtp = mocker.patch( - "airflow.providers.smtp.hooks.smtp.aiosmtplib.SMTP", autospec=True, side_effect=clients - ) + clients = [self._create_fake_async_smtp(Mock()) for _ in range(2)] + mock_smtp = mocker.patch("airflow.providers.smtp.hooks.smtp.aiosmtplib.SMTP", side_effect=clients) if quit_fails: clients[0].quit.side_effect = OSError("Quit failed") hook = SmtpHook() diff --git a/providers/smtp/tests/unit/smtp/notifications/test_smtp.py b/providers/smtp/tests/unit/smtp/notifications/test_smtp.py index 305534bf00702..02e8734539c44 100644 --- a/providers/smtp/tests/unit/smtp/notifications/test_smtp.py +++ b/providers/smtp/tests/unit/smtp/notifications/test_smtp.py @@ -26,7 +26,6 @@ import pytest -from airflow.providers.common.compat.sdk import Connection from airflow.providers.smtp.notifications.smtp import SmtpNotifier, send_smtp_notification from tests_common.test_utils.config import conf_vars @@ -98,15 +97,17 @@ def rendered(self) -> str: class TestSmtpNotifier: @mock.patch("smtplib.SMTP_SSL", autospec=True) - def test_repeated_notifications_reconnect(self, mock_smtp_ssl, create_connection_without_db): - create_connection_without_db( - Connection( - conn_id=SMTP_CONN_ID, - conn_type="smtp", - host="smtp.example.com", - port=465, - extra=json.dumps({"disable_tls": True}), - ) + def test_repeated_notifications_reconnect(self, mock_smtp_ssl, monkeypatch): + monkeypatch.setenv( + f"AIRFLOW_CONN_{SMTP_CONN_ID.upper()}", + json.dumps( + { + "conn_type": "smtp", + "host": "smtp.example.com", + "port": 465, + "extra": {"disable_tls": True}, + } + ), ) clients = [mock.create_autospec(smtplib.SMTP, instance=True) for _ in range(2)] mock_smtp_ssl.side_effect = clients