dify/api/tests/unit_tests/libs/test_smtp_client.py at 6facade85c2723292b4d8585dfa587921059a5b7 · gitarcode/dify · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
from unittest.mock import ANY, MagicMock, patch

import pytest

from libs.smtp import SMTPClient


def _mail() -> dict:
    return {"to": "user@example.com", "subject": "Hi", "html": "<b>Hi</b>"}


@patch("libs.smtp.smtplib.SMTP")
def test_smtp_plain_success(mock_smtp_cls: MagicMock):
    mock_smtp = MagicMock()
    mock_smtp_cls.return_value = mock_smtp

    client = SMTPClient(server="smtp.example.com", port=25, username="", password="", _from="noreply@example.com")
    client.send(_mail())

    mock_smtp_cls.assert_called_once_with("smtp.example.com", 25, timeout=10, local_hostname=ANY)
    mock_smtp.sendmail.assert_called_once()
    mock_smtp.quit.assert_called_once()


@patch("libs.smtp.smtplib.SMTP")
def test_smtp_tls_opportunistic_success(mock_smtp_cls: MagicMock):
    mock_smtp = MagicMock()
    mock_smtp_cls.return_value = mock_smtp

    client = SMTPClient(
        server="smtp.example.com",
        port=587,
        username="user",
        password="pass",
        _from="noreply@example.com",
        use_tls=True,
        opportunistic_tls=True,
    )
    client.send(_mail())

    mock_smtp_cls.assert_called_once_with("smtp.example.com", 587, timeout=10, local_hostname=ANY)
    assert mock_smtp.ehlo.call_count == 2
    mock_smtp.starttls.assert_called_once()
    mock_smtp.login.assert_called_once_with("user", "pass")
    mock_smtp.sendmail.assert_called_once()
    mock_smtp.quit.assert_called_once()


@patch("libs.smtp.smtplib.SMTP_SSL")
def test_smtp_tls_ssl_branch_and_timeout(mock_smtp_ssl_cls: MagicMock):
    # Cover SMTP_SSL branch and TimeoutError handling
    mock_smtp = MagicMock()
    mock_smtp.sendmail.side_effect = TimeoutError("timeout")
    mock_smtp_ssl_cls.return_value = mock_smtp

    client = SMTPClient(
        server="smtp.example.com",
        port=465,
        username="",
        password="",
        _from="noreply@example.com",
        use_tls=True,
        opportunistic_tls=False,
    )
    with pytest.raises(TimeoutError):
        client.send(_mail())
    mock_smtp.quit.assert_called_once()


@patch("libs.smtp.smtplib.SMTP")
def test_smtp_generic_exception_propagates(mock_smtp_cls: MagicMock):
    mock_smtp = MagicMock()
    mock_smtp.sendmail.side_effect = RuntimeError("oops")
    mock_smtp_cls.return_value = mock_smtp

    client = SMTPClient(server="smtp.example.com", port=25, username="", password="", _from="noreply@example.com")
    with pytest.raises(RuntimeError):
        client.send(_mail())
    mock_smtp.quit.assert_called_once()


@patch("libs.smtp.smtplib.SMTP")
def test_smtp_smtplib_exception_in_login(mock_smtp_cls: MagicMock):
    # Ensure we hit the specific SMTPException except branch
    import smtplib

    mock_smtp = MagicMock()
    mock_smtp.login.side_effect = smtplib.SMTPException("login-fail")
    mock_smtp_cls.return_value = mock_smtp

    client = SMTPClient(
        server="smtp.example.com",
        port=25,
        username="user",  # non-empty to trigger login
        password="pass",
        _from="noreply@example.com",
    )
    with pytest.raises(smtplib.SMTPException):
        client.send(_mail())
    mock_smtp.quit.assert_called_once()