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(predictor, images, expected_predictions):
    predictions = list(predictor.predict_images(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)