119 lines
3 KiB
Python
119 lines
3 KiB
Python
from typing import Any, Iterable
|
|
|
|
import pytest
|
|
from great_ai.utilities import WorkerException, threaded_parallel_map
|
|
from typing_extensions import Never
|
|
|
|
COUNT = int(1e5) + 3
|
|
|
|
|
|
def test_simple_case() -> None:
|
|
assert list(
|
|
threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=4)
|
|
) == [v**2 for v in range(COUNT)]
|
|
|
|
|
|
def test_with_iterable() -> None:
|
|
from time import sleep
|
|
|
|
def my_generator() -> Iterable[int]:
|
|
for i in range(10):
|
|
yield i
|
|
sleep(0.1)
|
|
|
|
expected = [v**3 for v in range(10)]
|
|
|
|
assert (
|
|
list(threaded_parallel_map(lambda x: x**3, my_generator(), chunk_size=1))
|
|
== expected
|
|
)
|
|
|
|
|
|
def test_simple_case_invalid_values() -> None:
|
|
with pytest.raises(AssertionError):
|
|
list(threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=0))
|
|
|
|
with pytest.raises(AssertionError):
|
|
list(threaded_parallel_map(lambda v: v**2, range(COUNT), chunk_size=0))
|
|
|
|
|
|
def test_this_worker_exception() -> None:
|
|
def my_generator() -> Iterable[int]:
|
|
yield 1
|
|
yield 2
|
|
yield 3
|
|
assert False
|
|
|
|
with pytest.raises(AssertionError):
|
|
list(
|
|
threaded_parallel_map(
|
|
lambda v: v**2, my_generator(), concurrency=2, chunk_size=2
|
|
)
|
|
)
|
|
|
|
with pytest.raises(AssertionError):
|
|
list(
|
|
threaded_parallel_map(
|
|
lambda v: v**2, my_generator(), concurrency=1, chunk_size=2
|
|
)
|
|
)
|
|
|
|
|
|
def test_ignore_this_worker_exception() -> None:
|
|
def my_generator() -> Iterable[float]:
|
|
yield 1
|
|
yield 2
|
|
yield 3
|
|
yield 1 / 0
|
|
|
|
assert list(
|
|
threaded_parallel_map(
|
|
lambda v: v**2,
|
|
my_generator(),
|
|
concurrency=2,
|
|
chunk_size=2,
|
|
ignore_exceptions=True,
|
|
)
|
|
) == [
|
|
1,
|
|
4,
|
|
] # the second chunk is ruined because of the error
|
|
|
|
|
|
def test_worker_worker_exception() -> None:
|
|
def oh_no(_: Any) -> Never:
|
|
raise ValueError("hi")
|
|
|
|
with pytest.raises(WorkerException):
|
|
list(threaded_parallel_map(oh_no, range(COUNT), concurrency=2))
|
|
|
|
with pytest.raises(WorkerException):
|
|
list(threaded_parallel_map(oh_no, range(COUNT), concurrency=1))
|
|
|
|
|
|
def test_ignore_worker_worker_exception() -> None:
|
|
def oh_no(_: Any) -> Never:
|
|
raise ValueError("hi")
|
|
|
|
assert (
|
|
list(
|
|
threaded_parallel_map(
|
|
oh_no, range(3), concurrency=2, ignore_exceptions=True
|
|
)
|
|
)
|
|
== [None] * 3
|
|
)
|
|
assert (
|
|
list(
|
|
threaded_parallel_map(
|
|
oh_no, range(3), concurrency=1, ignore_exceptions=True
|
|
)
|
|
)
|
|
== [None] * 3
|
|
)
|
|
|
|
|
|
def test_no_op() -> None:
|
|
assert list(threaded_parallel_map(lambda v: v**2, [])) == []
|
|
assert list(threaded_parallel_map(lambda v: v**2, [], chunk_size=100)) == []
|
|
assert list(threaded_parallel_map(lambda v: v**2, [], concurrency=100)) == []
|