如何对Python中含无限while循环的代码进行单元测试？

要测试类A的exe方法，同时终止infi里的无限循环，核心思路是给循环加可控制的终止信号，下面是具体实现步骤：

1. 修改原类，添加终止机制 #

原代码里的infi用while True硬编码了无限循环，我们需要给它加一个停止开关，比如用threading.Event作为信号：

import threading
import time

class A:
    def __init__(self):
        self.stop_event = threading.Event()  # 初始化停止信号

    def exe(self):
        self.x = threading.Event()
        # 注意：原代码中线程name参数不能传类A，要改成字符串
        self.thread = threading.Thread(name="A_infi_thread", target=self.infi, args=(self.x,), daemon=True)
        self.thread.start()

    def infi(self, x):
        # 把while True改成检查停止信号
        while not self.stop_event.is_set():
            # 原业务逻辑保持不变
            start = True  
            if start:
                start = False
                x.set()
            time.sleep(5)

    def stop(self):
        # 触发停止信号，并等待线程安全退出
        self.stop_event.set()
        self.thread.join()

2. 编写测试代码 #

测试时，先调用exe方法，验证功能后再触发stop终止循环：

def test_a_exe():
    # 初始化实例并启动线程
    a = A()
    a.exe()
    
    # 验证exe方法的核心逻辑：x是否被正确设置
    assert a.x.is_set() is True
    
    # 终止循环并等待线程结束
    a.stop()
    
    # 确认线程已经退出
    assert a.thread.is_alive() is False

if __name__ == "__main__":
    test_a_exe()

3. 额外技巧：测试时缩短等待时间 #

原代码里time.sleep(5)会让测试变慢，测试阶段可以临时把sleep时间改短，或者用mock替换time.sleep，比如用unittest.mock.patch：

from unittest.mock import patch

def test_a_exe_fast():
    a = A()
    # 替换sleep为立即返回，加快测试速度
    with patch("time.sleep", return_value=None):
        a.exe()
        assert a.x.is_set() is True
        a.stop()
        assert a.thread.is_alive() is False

其他可选方案（不修改原类） #

如果不想改动原类代码，可以在测试时用猴子补丁替换infi方法，跳过无限循环：

def test_a_exe_with_patch():
    a = A()
    # 临时替换infi方法，只执行核心逻辑不进入循环
    original_infi = a.infi
    def mock_infi(x):
        start = True
        if start:
            start = False
            x.set()
    a.infi = mock_infi
    
    a.exe()
    assert a.x.is_set() is True

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