用yield为函数编写unittest

Question

我试图为使用生成器的函数编写一个unittest。以下是我的代码：

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

我的单元测试代码：

^{pr2}$

如果extract_data函数使用return而不是yield，则此测试可以顺利通过。如何编写生成器的测试？在

Answer 1

我知道我需要做什么。我需要把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()

Answer 2

您的代码，稍作修改，不需要unittest：

import re

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():
    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)
    return expected_res == res

print(test_extract_data())

这将打印False

问题是，当您执行return时，在您的例子中，函数返回一个str。然而，{cd4{cd6>返回 ^{pr2}$

这将打印True。在

为了说明，在Python command prompt：

>>> type("hello")
<class 'str'>
>>> def gen():
...     yield "hello"
... 
>>> type(gen())
<class 'generator'>

您的另一个选择（可能更好，取决于您的用例）是通过将generator对象的结果转换为list或{}来测试generator的所有结果是否正确，然后比较是否相等：

import re

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():
    sample_input = ['<tr><h1>keyThis</h1><h2>\xc2</h2><h3>\xa0</h3><h4>view\xc2\xa0book\xc2\xa0info</h4><h5>Test Passes</h5></tr>', '<tr><h1>keyThis</h1><h2>\xc2</h2><h3>\xa0</h3><h4>view\xc2\xa0book\xc2\xa0info</h4><h5>Test Passes Too!</h5></tr>']
    expected_res = ['This Test Passes', 'This Test Passes Too!']
    res = extract_data(sample_input)
    return expected_res == list(res)

print(test_extract_data())