如何使用 pytest 测试迭代器函数？

Question

我有一个函数，它使用批量大小参数在可迭代对象上提供批量迭代器：

def batch_iterator(iterable: Iterable[Row], batch_size: int) -> Iterator:
    """
    Slices off a batch of values from an iterable, and returns it as an iterator.
    """
    return iter(lambda: list(itertools.islice(iterable, batch_size)), [])

我想用pytest测试一下这个函数。这是我尝试过的错误：

def test_batch_iterator():

    words_tuple = ('jimi', 'bertrand', 'alain', 'buck', 'apple', 'banana', 'cherry', 'oak', 'maple', 'python')
    result = batch_iterator(3, words_tuple)
    assert result == iter(words_tuple[:2])

#----------------------------------------------------------

the pytest result at the console:
...
Expected :<tuple_iterator object at 0x00000216F2F23130>
Actual   :<callable_iterator object at 0x00000216F2F23070>
<Click to see difference>

def test_batch_iterator():
    
        words_tuple = ('jimi', 'bertrand', 'alain', 'buck', 'apple', 'banana', 'cherry', 'oak', 'maple', 'python')
        result = batch_iterator(3, words_tuple)
>       assert result == iter(words_tuple[:2])
E       assert <callable_iterator object at 0x00000216F2F23070> == <tuple_iterator object at 0x00000216F2F23130>

Answer 1

batch_iterator(words_tuple, 3) returns 带有 ['jimi', 'bertrand', 'alain'] 的迭代器。要从中获取列表，您可以使用 next。还需要在调用batch_iterator()时切换参数，将words_tuple改为列表，或者在assertion

中强制转换

words_tuple = ['jimi', 'bertrand', 'alain', 'buck', 'apple', 'banana', 'cherry', 'oak', 'maple', 'python']
result = batch_iterator(words_tuple, 3)
assert next(result) == words_tuple[:3]

更新

目前 batch_iterator 创建一个仅包含前 3 个值的 iterator。如果你想要 words_tuple 中的所有条目，你可以使用类似

的东西

def batch_iterator(iterable: Iterable[Row], batch_size: int) -> Iterator:
    """
    Slices off a batch of values from an iterable, and returns it as an iterator.
    """
    return iter(iterable[i:i+batch_size] for i in range(0, len(iterable), batch_size))

words_tuple = ('jimi', 'bertrand', 'alain', 'buck', 'apple', 'banana', 'cherry', 'oak', 'maple', 'python')
result = batch_iterator(words_tuple, 3)
assert next(result) == words_tuple[:3]

如何使用 pytest 测试迭代器函数？

How to test an iterator function with pytest?

python

unit-testing

pytest