36 lines
1019 B
Python
36 lines
1019 B
Python
import pytest
|
|
|
|
from image_prediction.utils import chunk_iterable
|
|
|
|
|
|
@pytest.mark.parametrize("estimator_type", ["mock", "keras"])
|
|
@pytest.mark.parametrize("batch_size", [0, 1, 2, 16, 32, 64])
|
|
def test_predict(image_classifier, images, expected_predictions):
|
|
predictions = list(image_classifier.predict(images))
|
|
assert predictions == expected_predictions
|
|
|
|
|
|
def test_chunk_iterable_exact_split():
|
|
a, b = chunk_iterable(range(10), chunk_size=5)
|
|
assert a == tuple(range(5))
|
|
assert b == tuple(range(5, 10))
|
|
|
|
|
|
def test_chunk_iterable_no_split():
|
|
a = next(chunk_iterable(range(10), chunk_size=10))
|
|
assert a == tuple(range(10))
|
|
|
|
|
|
def test_chunk_iterable_last_partial():
|
|
a, b, c, d = chunk_iterable(range(10), chunk_size=3)
|
|
assert d == (9,)
|
|
|
|
|
|
def test_chunk_iterable_empty():
|
|
with pytest.raises(StopIteration):
|
|
next(chunk_iterable(range(0), chunk_size=3))
|
|
|
|
|
|
def test_chunk_iterable_less_than_chunk_size_elements():
|
|
assert next(chunk_iterable(range(2), chunk_size=5)) == (0, 1)
|