30
30
from airflow .providers .smtp .hooks .smtp import SmtpHook
31
31
from airflow .utils import db
32
32
from airflow .utils .session import create_session
33
+ from tests .test_utils .config import conf_vars
33
34
34
35
smtplib_string = "airflow.providers.smtp.hooks.smtp.smtplib"
35
36
@@ -75,13 +76,16 @@ def setup_method(self):
75
76
)
76
77
77
78
@patch (smtplib_string )
78
- def test_connect_and_disconnect (self , mock_smtplib ):
79
+ @patch ("ssl.create_default_context" )
80
+ def test_connect_and_disconnect (self , create_default_context , mock_smtplib ):
79
81
mock_conn = _create_fake_smtp (mock_smtplib )
80
82
81
83
with SmtpHook ():
82
84
pass
83
-
84
- mock_smtplib .SMTP_SSL .assert_called_once_with (host = "smtp_server_address" , port = 465 , timeout = 30 )
85
+ assert create_default_context .called
86
+ mock_smtplib .SMTP_SSL .assert_called_once_with (
87
+ host = "smtp_server_address" , port = 465 , timeout = 30 , context = create_default_context .return_value
88
+ )
85
89
mock_conn .login .assert_called_once_with ("smtp_user" , "smtp_password" )
86
90
assert mock_conn .close .call_count == 1
87
91
@@ -201,12 +205,90 @@ def test_hook_conn(self, mock_smtplib, mock_hook_conn):
201
205
202
206
@patch ("smtplib.SMTP_SSL" )
203
207
@patch ("smtplib.SMTP" )
204
- def test_send_mime_ssl (self , mock_smtp , mock_smtp_ssl ):
208
+ @patch ("ssl.create_default_context" )
209
+ def test_send_mime_ssl (self , create_default_context , mock_smtp , mock_smtp_ssl ):
205
210
mock_smtp_ssl .return_value = Mock ()
206
211
with SmtpHook () as smtp_hook :
207
212
smtp_hook .send_email_smtp (to = "to" , subject = "subject" , html_content = "content" , from_email = "from" )
208
213
assert not mock_smtp .called
209
- mock_smtp_ssl .assert_called_once_with (host = "smtp_server_address" , port = 465 , timeout = 30 )
214
+ assert create_default_context .called
215
+ mock_smtp_ssl .assert_called_once_with (
216
+ host = "smtp_server_address" , port = 465 , timeout = 30 , context = create_default_context .return_value
217
+ )
218
+
219
+ @patch ("smtplib.SMTP_SSL" )
220
+ @patch ("smtplib.SMTP" )
221
+ @patch ("ssl.create_default_context" )
222
+ def test_send_mime_ssl_none_email_context (self , create_default_context , mock_smtp , mock_smtp_ssl ):
223
+ mock_smtp_ssl .return_value = Mock ()
224
+ with conf_vars ({("smtp" , "smtp_ssl" ): "True" , ("email" , "ssl_context" ): "none" }):
225
+ with SmtpHook () as smtp_hook :
226
+ smtp_hook .send_email_smtp (
227
+ to = "to" , subject = "subject" , html_content = "content" , from_email = "from"
228
+ )
229
+ assert not mock_smtp .called
230
+ assert not create_default_context .called
231
+ mock_smtp_ssl .assert_called_once_with (host = "smtp_server_address" , port = 465 , timeout = 30 , context = None )
232
+
233
+ @patch ("smtplib.SMTP_SSL" )
234
+ @patch ("smtplib.SMTP" )
235
+ @patch ("ssl.create_default_context" )
236
+ def test_send_mime_ssl_none_smtp_provider_context (self , create_default_context , mock_smtp , mock_smtp_ssl ):
237
+ mock_smtp_ssl .return_value = Mock ()
238
+ with conf_vars ({("smtp" , "smtp_ssl" ): "True" , ("smtp_provider" , "ssl_context" ): "none" }):
239
+ with SmtpHook () as smtp_hook :
240
+ smtp_hook .send_email_smtp (
241
+ to = "to" , subject = "subject" , html_content = "content" , from_email = "from"
242
+ )
243
+ assert not mock_smtp .called
244
+ assert not create_default_context .called
245
+ mock_smtp_ssl .assert_called_once_with (host = "smtp_server_address" , port = 465 , timeout = 30 , context = None )
246
+
247
+ @patch ("smtplib.SMTP_SSL" )
248
+ @patch ("smtplib.SMTP" )
249
+ @patch ("ssl.create_default_context" )
250
+ def test_send_mime_ssl_none_smtp_provider_default_email_context (
251
+ self , create_default_context , mock_smtp , mock_smtp_ssl
252
+ ):
253
+ mock_smtp_ssl .return_value = Mock ()
254
+ with conf_vars (
255
+ {
256
+ ("smtp" , "smtp_ssl" ): "True" ,
257
+ ("email" , "ssl_context" ): "default" ,
258
+ ("smtp_provider" , "ssl_context" ): "none" ,
259
+ }
260
+ ):
261
+ with SmtpHook () as smtp_hook :
262
+ smtp_hook .send_email_smtp (
263
+ to = "to" , subject = "subject" , html_content = "content" , from_email = "from"
264
+ )
265
+ assert not mock_smtp .called
266
+ assert not create_default_context .called
267
+ mock_smtp_ssl .assert_called_once_with (host = "smtp_server_address" , port = 465 , timeout = 30 , context = None )
268
+
269
+ @patch ("smtplib.SMTP_SSL" )
270
+ @patch ("smtplib.SMTP" )
271
+ @patch ("ssl.create_default_context" )
272
+ def test_send_mime_ssl_default_smtp_provider_none_email_context (
273
+ self , create_default_context , mock_smtp , mock_smtp_ssl
274
+ ):
275
+ mock_smtp_ssl .return_value = Mock ()
276
+ with conf_vars (
277
+ {
278
+ ("smtp" , "smtp_ssl" ): "True" ,
279
+ ("email" , "ssl_context" ): "none" ,
280
+ ("smtp_provider" , "ssl_context" ): "default" ,
281
+ }
282
+ ):
283
+ with SmtpHook () as smtp_hook :
284
+ smtp_hook .send_email_smtp (
285
+ to = "to" , subject = "subject" , html_content = "content" , from_email = "from"
286
+ )
287
+ assert not mock_smtp .called
288
+ assert create_default_context .called
289
+ mock_smtp_ssl .assert_called_once_with (
290
+ host = "smtp_server_address" , port = 465 , timeout = 30 , context = create_default_context .return_value
291
+ )
210
292
211
293
@patch ("smtplib.SMTP_SSL" )
212
294
@patch ("smtplib.SMTP" )
@@ -269,7 +351,10 @@ def test_send_mime_partial_failure(self, mock_smtp_ssl, mime_message_mock):
269
351
270
352
@patch ("airflow.models.connection.Connection" )
271
353
@patch ("smtplib.SMTP_SSL" )
272
- def test_send_mime_custom_timeout_retrylimit (self , mock_smtp_ssl , connection_mock ):
354
+ @patch ("ssl.create_default_context" )
355
+ def test_send_mime_custom_timeout_retrylimit (
356
+ self , create_default_context , mock_smtp_ssl , connection_mock
357
+ ):
273
358
mock_smtp_ssl ().sendmail .side_effect = smtplib .SMTPServerDisconnected ()
274
359
custom_retry_limit = 10
275
360
custom_timeout = 60
@@ -287,6 +372,10 @@ def test_send_mime_custom_timeout_retrylimit(self, mock_smtp_ssl, connection_moc
287
372
with pytest .raises (smtplib .SMTPServerDisconnected ):
288
373
smtp_hook .send_email_smtp (to = "to" , subject = "subject" , html_content = "content" )
289
374
mock_smtp_ssl .assert_any_call (
290
- host = fake_conn .host , port = fake_conn .port , timeout = fake_conn .extra_dejson ["timeout" ]
375
+ host = fake_conn .host ,
376
+ port = fake_conn .port ,
377
+ timeout = fake_conn .extra_dejson ["timeout" ],
378
+ context = create_default_context .return_value ,
291
379
)
380
+ assert create_default_context .called
292
381
assert mock_smtp_ssl ().sendmail .call_count == 10
0 commit comments