如何通过模拟库模拟套接字对象

Question

如何通过模拟库(在Python 3的情况下)从Python的标准库中模拟TCPSocket包装器？socketunittest.mock

这是我的包装:

import socket

import utils


class TCPSocket:
    def __init__(self):
        self.buf = ''

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self, host, port):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((host, port))

    def close(self):
        self.sock.close()

    def send(self, data):
        self.sock.send(data)

    def sendall(self, data):
        self.sock.sendall(data)

    # For best match with hardware and network realities,
    # the value of limit should be a relatively small power of 2, for example, 4096
    def recv_some(self, limit=4096):
        return self.sock.recv(limit)

    def recv_bytes(self, count):
        while len(self.buf) < count:
            self.buf += self.recv_some()
        parts = utils.split_first_n(self.buf, count)
        self.buf = parts[1]
        return parts[0]

    def recv_until(self, delim):
        while delim not in self.buf:
            self.buf += self.recv_some()
        parts = self.buf.split(delim, maxsplit=1)
        self.buf = parts[1]
        return parts[0]

我想测试recv_until和recv_bytes之类的函数是否真正需要它们.

Answer 1

它可以很简单

import mock   # or from unittest import mock


mock_socket = mock.Mock()
mock_socket.recv.return_value = data

然后使用mock_socket你将使用真正的套接字的地方.您还可以模拟创建套接字的任何内容,以返回类似于此处配置的模拟值,具体取决于您的需要.

对于您的情况,您可以进行模拟,socket.socket以便返回可以配置其方法的内容.请注意,mock_socket在此示例中是一个返回Socket对象的函数,而不是Socket对象本身.

with mock.patch('socket.socket') as mock_socket:
    mock_socket.return_value.recv.return_value = some_data
    t = TCPSocket()
    t.connect('example.com', 12345)  # t.sock is a mock object, not a Socket
self.assertEqual(t.recv_bytes(), whatever_you_expect)
t.sock.connect.assert_called_with(('example.com', 12345))