def test_pmap_performance(chunk_size, include_power):
def sleep_exp(x, power=2):
time.sleep(PERFORMANCE_SLEEP_TIME)
return x**power
number_partitions = len(
list(pbatch.partition(PERFORMANCE_ITEMS, chunk_size=chunk_size)))
# plus one for buffer of execution overhead (assuming overhead is
# less than PERFORMANCE_SLEEP_TIME)
max_time = PERFORMANCE_SLEEP_TIME * (number_partitions + 1)
if include_power:
powers = [3] * len(PERFORMANCE_ITEMS)
results, duration = time_list(pbatch.pmap,
sleep_exp,
PERFORMANCE_ITEMS,
powers,
chunk_size=chunk_size)
else:
results, duration = time_list(pbatch.pmap,
sleep_exp,
PERFORMANCE_ITEMS,
chunk_size=chunk_size)
if include_power:
assert results == [0, 1, 8, 27, 64, 125, 216, 343, 512, 729]
else:
assert results == [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
assert duration < max_time
def test_set_items(items, chunk_size, expected_list):
# cannot use dict items in a set (not hashable)
if isinstance(next(iter(items), None), dict):
return
partitions = pbatch.partition(items=set(items), chunk_size=chunk_size)
assert sorted(itertools.chain(*partitions)) == sorted(itertools.chain(*expected_list))
def test_consumes_iterator():
iterator = iter([1, 2, 3, 4])
partitions = pbatch.partition(iterator, 2)
assert next(partitions) == [1, 2]
assert list(iterator) == [3, 4]
assert list(partitions) == []
def test_generator_returned(items, chunk_size, expected_list):
partitions = pbatch.partition(items, chunk_size)
assert isinstance(partitions, GeneratorType)
if expected_list:
assert next(partitions) == expected_list[0]
for actual, expected in zip(partitions, expected_list[1:]):
assert actual == expected
else:
assert list(partitions) == []
def test_generator_items(items, chunk_size, expected_list):
def items_generator():
yield from items
assert list(pbatch.partition(items=items_generator(), chunk_size=chunk_size)) == expected_list
def test_iter_items(items, chunk_size, expected_list):
assert list(pbatch.partition(items=iter(items), chunk_size=chunk_size)) == expected_list
def test_items_and_chunk_size_kwargs(items, chunk_size, expected_list):
assert list(pbatch.partition(items=items, chunk_size=chunk_size)) == expected_list
def test_partitions(items, chunk_size, expected_list):
assert list(pbatch.partition(items, chunk_size)) == expected_list
def test_non_int_chunk_size(chunk_size):
with pytest.raises(AssertionError) as info:
list(pbatch.partition([], chunk_size))
assert str(info.value) == "Chunk size must be a positive int (or None)"
def test_invalid_arguments(args):
with pytest.raises(TypeError):
pbatch.partition(*args)