斷言多個方法被調用，只有構造函數有效

Question

我正在嘗試測試以下 function

def send_mail(config, message, raw_object):
    smtp_config = config['handlers']['smtp']

    session = smtplib.SMTP(smtp_config['host'], smtp_config['port'])
    if smtp_config['tls']:
        session.starttls()
    session.login(smtp_config['from'], smtp_config['password'])

    for to in smtp_config['to']:
        mail = MIMEMultipart()
        mail['From'] = smtp_config['from']
        mail['To'] = to
        mail['Subject'] = message

        body = yaml.safe_dump(raw_object)

        mail.attach(MIMEText(body, 'plain'))

        try:
            session.sendmail(smtp_config['from'], to, mail.as_string())
            logging.info(f"Handler:SMTP {to}: {message}")
        except smtplib.SMTPException as exc:
            logging.error("SMTPException:")
            logging.error(exc)

    session.quit()

我有以下測試

from unittest import TestCase
from unittest.mock import patch

from kubewatcher.handlers import send_mail


class Test(TestCase):
    @patch("smtplib.SMTP")
    def test_handle__send_mail(self, smtp):
        from_ = "from"
        password = "password"
        host = "host"
        port = 587
        tls = True
        to = ["to"]

        config = {
            "handlers": {
                "smtp": {
                    "from": from_,
                    "password": password,
                    "host": host,
                    "port": port,
                    "tls": tls,
                    "to": to
                }
            }
        }

        message = "message"

        raw_object = {}

        send_mail(config, message, raw_object)

        smtp.assert_called_once_with(host, port)
        smtp.starttls.assert_called_once()
        smtp.login.assert_called_once_with(from_, password)

第一個斷言smtp.assert_called_once_with(host, port)工作得很好。 但是整個測試失敗並出現以下錯誤

...
AssertionError: Expected 'starttls' to have been called once. Called 0 times.

Answer 1

這是您正在測試的代碼：

    session = smtplib.SMTP(smtp_config['host'], smtp_config['port'])
    if smtp_config['tls']:
        session.starttls()
    session.login(smtp_config['from'], smtp_config['password'])

此測試有效：

smtp = patch("smtplib.SMTP")  # Sorta; this is just shorthand
smtp.assert_called_once_with(host, port)

問題在這里：

smtp.starttls.assert_called_once()

但實際上失敗是正確的。 您的代碼不是調用smtplib.SMTP.starttls ，而是session.starttls ，這是smtplib.SMTP返回的東西

您可以通過以下方式解決此問題：

from unittest.mock import patch, Mock

mock_session = Mock()  # This is the thing we'll be inspecting later
smtp.return_value = mock_session  # `smtplib.SMTP` will return this object

send_mail(config, message, raw_object)

smtp.assert_called_once_with(host, port)
mock_session.starttls.assert_called_once()
mock_session.login.assert_called_once_with(from_, password)