如何使用pytest测试自定义批量迭代器函数？

问题原因梳理 #

• 参数传参顺序错误：原函数定义的第一个入参为可迭代对象，第二个为批量大小，测试代码中将3和words_tuple的传入顺序写反
• 不可直接比较迭代器对象：不同实现逻辑的迭代器实例就算产出内容完全一致，直接用==比较也会返回False，需要先把迭代器的所有输出消费转为列表这类可直接比较的序列
• 预期结果逻辑错误：batch_iterator的功能是按指定大小拆分迭代器返回批量，不是返回前N个元素的迭代器，batch_size=3时10个元素的输入会产出4个批量

正确测试用例示例 #

首先确保导入必要依赖：

import itertools
from typing import Iterable, Iterator

基础功能测试代码：

def batch_iterator(iterable: Iterable, 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)), [])

def test_batch_iterator_normal_scenario():
    words_tuple = ('jimi', 'bertrand', 'alain', 'buck', 'apple', 'banana', 'cherry', 'oak', 'maple', 'python')
    # 严格按照函数定义顺序传入参数
    result = batch_iterator(words_tuple, 3)
    # 迭代器转为列表后再比较内容
    assert list(result) == [
        ['jimi', 'bertrand', 'alain'],
        ['buck', 'apple', 'banana'],
        ['cherry', 'oak', 'maple'],
        ['python']
    ]

你可以补充边界场景测试，覆盖更多异常情况：

def test_batch_iterator_batch_size_larger_than_length():
    # 批量大小大于元素总数的场景
    input_iter = [1,2,3]
    result = list(batch_iterator(input_iter, 5))
    assert result == [[1,2,3]]

def test_batch_iterator_exact_multiple():
    # 元素总数是批量大小整数倍的场景
    input_iter = [1,2,3,4,5,6]
    result = list(batch_iterator(input_iter, 2))
    assert result == [[1,2], [3,4], [5,6]]

def test_batch_iterator_empty_input():
    # 输入为空可迭代对象的场景
    input_iter = []
    result = list(batch_iterator(input_iter, 3))
    assert result == []

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