Skip to main content
 首页 » 编程设计

python之使用 yield 为函数编写单元测试

2025年05月04日28JeffreyZhao

我正在尝试为使用生成器的函数编写单元测试。下面是我的代码:

def extract_data(body): 
    for i in body: 
        a = re.sub('<[^<]+?>', '', str(i)) 
        b = re.sub('view\xc2\xa0book\xc2\xa0info', '', str(a)) 
        c = re.sub('key', '', str(b)) 
        d = re.sub('\xc2', ' ', str(c)) 
        e = re.sub('\xa0', '', str(d)) 
        yield e 

我的单元测试代码:

    def test_extract_data(self): 
        sample_input = ['<tr><h1>keyThis</h1><h2>\xc2</h2><h3>\xa0</h3><h4>view\xc2\xa0book\xc2\xa0info</h4><h5>Test Passes</h5></tr>'] 
        expected_res = 'This Test Passes' 
        res = extract_data(sample_input) 
 
        self.assertEqual(expected_res, res) 

如果 extract_data 函数使用 return 而不是 yield,则此测试会毫无问题地通过。如何为生成器编写测试?

请您参考如下方法:

我想出了我需要做什么。我需要将 res 放入列表中。就是这样。比我预期的要简单得多。所以这是现在的样子:

class TestScrapePage(unittest.TestCase): 
 
    def test_extract_data(self): 
        sample_input = ['<tr><h1>keyThis</h1><h2>\xc2</h2><h3>\xa0</h3><h4>view\xc2\xa0book\xc2\xa0info</h4><h5>Test Passes</h5></tr>'] 
        expected_res = ['This Test Passes'] 
        res = list(extract_data(sample_input)) 
 
    self.assertEqual(expected_res, res) 
 
if __name__ == '__main__': 
    unittest.main()