在Python中模拟套接字连接

Question

我正在尝试为Python中的一个类编写单元测试。该类在init时打开一个TCP套接字。我想对此进行模拟，以便可以断言使用正确值调用了连接，但在单元测试中实际上并没有发生。我已经尝试过使用MagicMock、patch等方法，但还没有找到解决方案。

到目前为止，我的类看起来像这样：

import socket

class MyClass(object):

    def __init__(self):
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.connect('0.0.0.0', '6767')

- CountCet

---

为什么不直接连接呢？ - Joran Beasley

因为连接需要花费太长时间来建立，如果不使用模拟数据可能会导致构建过程停滞。这就是我们需要使用模拟数据的原因。 - markroxor

Answer 1

如果您仅想断言connect被正确调用，那么简单的做法是

import mock
import socket

class MyClass(object):

    def __init__(self):
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.connect('0.0.0.0', '6767')

with mock.patch('socket.socket'):
    c = MyClass()
    c.tcp_socket.connect.assert_called_with('0.0.0.0', '6767')

如果你需要首先导入一个模块才能访问 MyClass，那么你需要稍微调整一下 patch：

from mymodule import MyClass
import mock

with mock.patch('mymodule.socket.socket'):
    c = MyClass()
    c.tcp_socket.connect.assert_called_with('0.0.0.0', '6767')

- chepner

---

在Python 3.3及以上版本中，你可以直接使用import unittest.mock来导入mock，因为它现在已经包含在标准库中了：https://docs.python.org/3/library/unittest.mock.html - phoenix