在 Python 中模拟复杂的数据结构

Question

我有一组旨在管理数据库集群的脚本。我已将这些脚本转换为面向对象的设计，并且在尝试编写单元测试时遇到了麻烦。这是我要实现的目标的简化示例：

class DBServer:
    def __init__(self,hostname):
        self.hostname=hostname

    def get_instance_status(self):
        if (<instance_is_online>):
            return True
        else:
            return false

    def shutdown_instance(self):
        #<insert code here to stop the DB instance on that node>


class DBCluster:
    def __init__(self,clustername):
        self.servers=[DBServer('hostname1'), DBServer('hostname2'), DBServer('hostname3'), DBServer('hostname4')]
    
    def get_servers_online(self):
        #returns an array of DBServer objects where the DB instance is online
        serversonline=[]
        for server in self.servers:
            if server.get_instance_status():
                serversonline.append(server)
        return serversonline
            
    def shutdown_cluster(self):
        for server in self.servers:
            if server.get_instance_status():
                server.shutdown_instance()

如您所见，DBServer 是一个具有主机名和一些可以在每个服务器上运行的功能的对象。 DBCluster 是一个包含 DBServer 数组的对象，以及一些管理整个集群的函数。

想象一个工作流，其中用户想要获取数据库实例在线的集群中的服务器列表。然后用户应该选择一个并将其关闭：

def shutdown_user_instance(node_number):
    #get list of online instances in cluster, display to user and shut down the chosen instance:
    cluster=DBCluster()
    servers=cluster.get_servers_online()
    print ('Choose server to shut down:')
    i=0
    for server in servers:
        print (str(i) +': ' +server.hostname)
        i++
    reply=str(input('Option: '))
    instance=int(reply)
    servers[instance].shutdown_instance()

我们如何对这个函数进行单元测试？我认为它应该模拟 cluster.get_instances_online() 函数以返回模拟 DBServer 对象的数组。然后我可以为该数组中的每个对象创建 server.hostname 的返回值，然后模拟 server.shutdown_instance()。问题是我不知道从哪里开始。像这样？

class TestCluster(unittest.TestCase):

    @patch(stack.DBCluster)
    @patch(stack.DBServer)
    @patch(builtins.input)
    @patch.object(stack.DBserver,'shutdown_instance')
    def test_shutdown(self, mock_cluster, mock_server, mock_input, mock_shutdown):
        mock_server.get_instance_status.side_effect[True,False,True,False] #for the purpose of this test, let's say nodes 0 and 2 are online, 1 and 3 are offline
        mock_cluster.get_servers_online.return_value=[mock_server,mock_server] #How should I achieve this array of mocked objects to test??
        mock_input.return_value=0
        shutdown_user_instance() #theoretically this should mock a shutdown of node 0

Answer 1

这是一个通过pytest的独立示例：

class DBCluster:
    def get_servers_online():
        pass

def shutdown_user_instance(node_number):
    #get list of online instances in cluster, display to user and shut down the chosen instance:
    cluster = DBCluster()
    servers=cluster.get_servers_online()
    print ('Choose server to shut down:')
    i=0
    for server in servers:
        print (str(i) +': ' +server.hostname)
        i += 1
    reply=str(input('Option: '))
    instance=int(reply)
    servers[instance].shutdown_instance()


from unittest.mock import patch, MagicMock

@patch.object(DBCluster, "__new__")  # same as @patch("stack.DBCluster")
@patch("builtins.input")
def test_shutdown(mock_input, mock_cluster):
    mock_servers = [MagicMock() for _ in range(3)]
    mock_cluster.return_value.get_servers_online.return_value = mock_servers
    mock_input.return_value = "1"
    shutdown_user_instance(0)
    assert [s.shutdown_instance.call_count for s in mock_servers] == [0, 1, 0]

关键是您只需要修补被测代码中直接使用的符号。一旦你用一个模拟对象修补了DBCluster，你就可以在你的测试中使用那个模拟和它的模拟属性。特别是，研究这条线并确保你完全理解它在做什么：

mock_cluster.return_value.get_servers_online.return_value = mock_servers

mock_cluster 是DBCluster 构造函数的模拟——所以它的return_value 是DBCluster() 将在被测代码中返回的值，即这是cluster 将具有的值。反过来，我们希望模拟对象的get_servers_online 有一个return_value，然后在被测代码中变为servers，并且这些元素中的每一个都是一个模拟，其shutdown_instance 属性可以在之后检查代码被执行：

assert [s.shutdown_instance.call_count for s in mock_servers] == [0, 1, 0]

断言我们关闭了实例 1，但没有关闭实例 0 或 2。（我总是使用 assert mock.call_count == 1 而不是 mock.assert_called()，因为如果你打错了前者，你的测试将会以明显的方式中断，而如果你打错了后者你的测试会在不应该通过的时候通过。）

还要特别注意patch/patch.object 语法——当您刚接触单元测试时，这是最容易出错的事情（通常也是最难调试的）。