如何在Python unittest中用单Socket连接测试TCP服务

你完全可以借助Python unittest框架的类级别初始化/清理方法，实现整个测试套件共享一个TCP Socket连接，不用每个测试用例都重复创建、连接、关闭socket。具体来说，用setUpClass()和tearDownClass()这两个方法——它们会分别在测试类的所有测试用例执行前、执行后各运行一次，刚好匹配你的需求。

下面是修改后的完整代码示例：

import socket
import unittest
from unittest import TestCase

class TestService(TestCase):
    # 定义类级别的socket对象，所有测试用例共享
    client_socket = None
    port_number = 12345  # 替换成你的服务实际端口号

    @classmethod
    def setUpClass(cls):
        """所有测试用例执行前仅运行一次：建立socket连接"""
        cls.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        cls.client_socket.connect(('', cls.port_number))
        print("全局Socket连接已建立")

    def test_1(self):
        self.client_socket.send('message1'.encode())
        response = self.client_socket.recv(1024).decode()
        self.assertEqual(response, 'reply1')

    def test_2(self):
        self.client_socket.send('message2'.encode())
        response = self.client_socket.recv(1024).decode()
        self.assertEqual(response, 'reply2')

    @classmethod
    def tearDownClass(cls):
        """所有测试用例执行完成后仅运行一次：关闭socket连接"""
        if cls.client_socket:
            cls.client_socket.close()
            print("全局Socket连接已关闭")

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

关键细节说明： #

• 用@classmethod装饰器标记setUpClass和tearDownClass：这两个方法属于测试类本身，而非单个测试实例，因此只会执行一次
• 把client_socket设为类属性：这样所有测试用例的实例都能访问同一个连接对象，实现共享
• 连接初始化与清理分离：setUpClass负责创建并连接socket，tearDownClass负责最后关闭，避免资源泄漏
• 测试用例简化：每个测试方法直接使用self.client_socket收发数据，不用再重复处理连接逻辑

额外提示： #

如果你的测试场景中，某个用例会导致服务端主动断开连接，你可以考虑在setUp()方法中添加连接状态检查，必要时重新建立连接——不过这会回到“每个测试前初始化”的模式。如果你的服务端支持长连接，上面的方案就足够高效且符合需求。

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