66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
import unittest
|
|
|
|
import pytest
|
|
|
|
from src.great_ai.utilities import parallel_map
|
|
|
|
COUNT = int(1e5) + 3
|
|
|
|
|
|
class TestParallelMap(unittest.TestCase):
|
|
def test_simple_case_with_progress_bar(self) -> None:
|
|
inputs = range(COUNT)
|
|
expected = [v**2 for v in range(COUNT)]
|
|
|
|
assert list(parallel_map(lambda v: v**2, inputs, concurrency=10)) == expected
|
|
|
|
def test_with_iterable(self) -> None:
|
|
from time import sleep
|
|
|
|
def my_generator():
|
|
for i in range(10):
|
|
yield i
|
|
sleep(0.1)
|
|
|
|
expected = [v**3 for v in range(10)]
|
|
|
|
assert (
|
|
list(parallel_map(lambda x: x**3, my_generator(), chunk_size=1))
|
|
== expected
|
|
)
|
|
|
|
def test_simple_case_without_progress_bar(self) -> None:
|
|
inputs = range(COUNT)
|
|
expected = [v**2 for v in range(COUNT)]
|
|
|
|
assert (
|
|
list(parallel_map(lambda v: v**2, inputs, disable_logging=True))
|
|
== expected
|
|
)
|
|
|
|
def test_simple_case_invalid_values(self) -> None:
|
|
inputs = range(COUNT)
|
|
|
|
with pytest.raises(AssertionError):
|
|
list(parallel_map(lambda v: v**2, inputs, concurrency=0))
|
|
|
|
with pytest.raises(AssertionError):
|
|
list(parallel_map(lambda v: v**2, inputs, chunk_size=0))
|
|
|
|
def test_no_op(self) -> None:
|
|
assert list(parallel_map(lambda v: v**2, [], disable_logging=True)) == []
|
|
|
|
assert (
|
|
list(
|
|
parallel_map(lambda v: v**2, [], disable_logging=True, chunk_size=100)
|
|
)
|
|
== []
|
|
)
|
|
assert (
|
|
list(
|
|
parallel_map(
|
|
lambda v: v**2, [], disable_logging=True, concurrency=100
|
|
)
|
|
)
|
|
== []
|
|
)
|