用yield来写函数unittest

时间:2016-06-21 23:56:48

标签: python unit-testing generator

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

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,则此测试没有问题。如何为发生器编写测试?

2 个答案:

答案 0 :(得分:4)

我想出了我需要做的事情。我需要把res变成一个列表。就是这样。比我想象的要简单得多。所以这就是现在的样子:

undefined

答案 1 :(得分:0)

您的代码略有改动,不需要单元测试:

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。但是,当您执行yield时,它会返回generator类型对象,其next()函数返回str。所以,例如:

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 == next(res)

print(test_extract_data())

这会打印True

为了说明,请Python command prompt

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

您的其他选项(可能更好,取决于您的使用案例),是通过将generator对象的结果转换为{{1}来测试generator的所有结果是否正确}或list,然后比较相等:

tuple