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

Question

如何通过mock库（在Python 3的情况下为unittest.mock）从Python标准库中模拟socket的TCPSocket包装？

这是我的包装：

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))