如何在Python中Mock Socket模块？UPnP多播项目单元测试需求

没问题，我来教你怎么用unittest的mock工具模拟socket.recvfrom调用，这样测试UPnP搜索功能时不用依赖真实网络，既高效又稳定。

假设你的msearch方法完整逻辑是发送M-SEARCH请求、循环接收响应直到超时，我们可以通过unittest.mock来替换socket的真实行为，预设测试数据。

1. 编写测试文件 #

在项目根目录创建test_mucassdp.py，内容如下：

import unittest
import socket
from unittest.mock import patch, Mock
from upnp.mucassdp import ssdpObj

class TestSSDPObj(unittest.TestCase):
    @patch('upnp.mucassdp.socket')
    def test_msearch_device_discovery(self, mock_socket):
        # 模拟socket实例
        mock_sock = Mock()
        mock_socket.return_value = mock_sock

        # 预设recvfrom的行为：第一次返回设备响应，第二次抛出超时（模拟实际循环退出逻辑）
        mock_sock.recvfrom.side_effect = [
            # 模拟收到一个UPnP设备的响应数据包
            (
                b'HTTP/1.1 200 OK\r\n'
                b'LOCATION: http://192.168.1.100:1900/device.xml\r\n'
                b'ST: upnp:rootdevice\r\n'
                b'SERVER: Linux/4.19.0 UPnP/1.0 Device/1.0\r\n\r\n',
                ('192.168.1.100', 1900)
            ),
            socket.timeout  # 第二次调用触发超时，让循环退出
        ]

        # 执行要测试的方法
        discovered_devices = ssdpObj().msearch()

        # 断言socket初始化正确
        mock_socket.assert_called_once_with(socket.AF_INET, socket.SOCK_DGRAM)

        # 断言M-SEARCH消息正确发送
        expected_search_msg = (
            b'M-SEARCH * HTTP/1.1\r\n'
            b'HOST:239.255.255.250:1900\r\n'
            b'ST:upnp:rootdevice\r\n'  # 替换成你实际的ST字段内容
            b'MAN: "ssdp:discover"\r\n'
            b'MX: 3\r\n\r\n'
        )
        mock_sock.sendto.assert_called_once_with(expected_search_msg, ('239.255.255.250', 1900))

        # 断言recvfrom被调用了两次（一次收响应，一次超时）
        self.assertEqual(mock_sock.recvfrom.call_count, 2)

        # 断言设备解析逻辑正确
        self.assertEqual(len(discovered_devices), 1)
        self.assertEqual(discovered_devices[0]['ip'], '192.168.1.100')
        self.assertEqual(discovered_devices[0]['location'], 'http://192.168.1.100:1900/device.xml')

if __name__ == '__main__':
    unittest.main()

2. 关键细节解释 #

• Patch路径要正确：我们patch的是upnp.mucassdp.socket，也就是你的业务代码中实际导入使用的socket模块，而不是全局的socket，这样才能精准替换掉目标代码里的socket行为。
• 模拟多次调用行为：用side_effect来模拟recvfrom的多次调用——第一次返回预设的设备响应，第二次抛出超时异常，完全匹配真实场景下的循环接收逻辑。
• 多维度断言：不仅验证返回结果，还检查socket初始化参数、发送的消息内容、调用次数，确保整个流程的每一步都符合预期。

如果你的msearch里还有socket选项设置（比如SO_REUSEADDR、超时配置），可以添加额外断言：

# 断言socket选项设置正确
mock_sock.setsockopt.assert_any_call(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
mock_sock.settimeout.assert_called_once_with(3)

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