Python单元测试：如何在上下文管理器中模拟SMTP方法抛异常

我来帮你搞定这个没法触发异常测试日志逻辑的问题！核心是你可能没正确给mock的SMTP实例设置异常抛出规则，下面结合常见场景给你具体方案：

先明确核心逻辑 #

因为SMTP是通过上下文管理器调用的，所以mock时不能只替换smtplib.SMTP类，还要拿到上下文管理器返回的server实例，再给它的send_message方法绑定异常抛出逻辑。

场景1：待测试代码直接使用SMTP上下文管理器 #

假设你的待测试代码大概是这样：

# siteidentity_proxy_monitoring.py
import smtplib
import logging

logger = logging.getLogger(__name__)

def send_monitoring_alert(msg):
    try:
        with smtplib.SMTP('smtp.server.com', 587) as smtp_server:
            smtp_server.send_message(msg)
    except Exception as e:
        logger.error(f"邮件发送失败: {str(e)}")

对应的单元测试可以这么写：

import unittest
from unittest.mock import patch, MagicMock
import logging
from siteidentity_proxy_monitoring import send_monitoring_alert

class TestEmailAlert(unittest.TestCase):
    def test_send_message_exception_triggers_log(self):
        # 捕获日志输出，方便后续验证
        with self.assertLogs(logger='siteidentity_proxy_monitoring', level='ERROR') as log_recorder:
            # 替换待测试代码里的SMTP类
            with patch('siteidentity_proxy_monitoring.smtplib.SMTP') as mock_smtp_cls:
                # 获取上下文管理器返回的SMTP server实例
                mock_smtp_server = mock_smtp_cls.return_value.__enter__.return_value
                # 给send_message设置异常抛出逻辑
                mock_smtp_server.send_message.side_effect = ConnectionError("SMTP服务器连接超时")
                
                # 调用待测试函数
                send_monitoring_alert(MagicMock())
                
                # 验证send_message确实被调用过
                mock_smtp_server.send_message.assert_called_once()
                # 验证日志是否正确记录了异常信息
                self.assertIn("邮件发送失败: SMTP服务器连接超时", log_recorder.output[0])

场景2：待测试代码封装了SMTP调用 #

如果你的代码是通过自定义类封装了SMTP操作，比如：

# siteidentity_proxy_monitoring.py
from smtplib import SMTP

class AlertSender:
    def send(self, msg):
        try:
            with SMTP('smtp.server.com', 587) as server:
                server.send_message(msg)
        except Exception as e:
            logging.error(f"发送失败: {e}")

那mock时要注意patch的路径是待测试代码中导入的SMTP，也就是siteidentity_proxy_monitoring.SMTP，而不是smtplib.SMTP，否则mock不会生效。

关键踩坑点提醒 #

• 不要忽略上下文管理器的__enter__方法：SMTP上下文管理器会返回server实例，必须通过mock_cls.return_value.__enter__.return_value拿到这个实例才能设置send_message的异常。
• 用assertLogs简化日志验证：这是Python3.4+自带的工具，不用自己手动捕获日志流，直接断言日志内容即可。
• 遵循mock的路径规则："在哪里用，就patch哪里"，而不是patch原模块的类，否则待测试代码会使用真实的SMTP类，mock无效。

如果你的代码有特殊的封装逻辑，可以把具体片段贴出来，我再帮你调整方案！

内容的提问来源于stack exchange，提问作者John Edge